diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml
index fddb52631..99fcf989a 100644
--- a/.JuliaFormatter.toml
+++ b/.JuliaFormatter.toml
@@ -1,2 +1,5 @@
+whitespace_in_kwargs = true
remove_extra_newlines = true
trailing_comma = true
+separate_kwargs_with_semicolon = true
+long_to_short_function_def = true
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 6159b6ccd..8506f4306 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -6,6 +6,7 @@ on:
- main
tags: ['*']
pull_request:
+ workflow_dispatch:
concurrency:
# Skip intermediate builds: always.
@@ -51,10 +52,12 @@ jobs:
version: '1'
- uses: julia-actions/julia-buildpkg@v1
- uses: julia-actions/julia-docdeploy@v1
+ with:
+ install-package: false
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
- run: |
- julia --project=docs -e '
+ julia --project=docs -t auto -e '
using Documenter: DocMeta, doctest
using IncompressibleNavierStokes
DocMeta.setdocmeta!(
diff --git a/.github/workflows/FormatCheck.yml b/.github/workflows/FormatCheck.yml
index e91759519..1f89190ce 100644
--- a/.github/workflows/FormatCheck.yml
+++ b/.github/workflows/FormatCheck.yml
@@ -5,6 +5,7 @@ on:
branches:
- "main"
pull_request:
+ workflow_dispatch:
jobs:
build:
diff --git a/.github/workflows/SpellCheck.yml b/.github/workflows/SpellCheck.yml
new file mode 100644
index 000000000..c344f6dda
--- /dev/null
+++ b/.github/workflows/SpellCheck.yml
@@ -0,0 +1,13 @@
+name: Spell Check
+
+on: [pull_request]
+
+jobs:
+ typos-check:
+ name: Spell Check with Typos
+ runs-on: ubuntu-latest
+ steps:
+ - name: Checkout Actions Repository
+ uses: actions/checkout@v4
+ - name: Check spelling
+ uses: crate-ci/typos@master
diff --git a/.github/workflows/TagBot.yml b/.github/workflows/TagBot.yml
index 6d2efc1ce..7667893fb 100644
--- a/.github/workflows/TagBot.yml
+++ b/.github/workflows/TagBot.yml
@@ -5,7 +5,22 @@ on:
types:
- created
workflow_dispatch:
-
+ inputs:
+ lookback:
+ default: 3
+permissions:
+ actions: read
+ checks: read
+ contents: write
+ deployments: read
+ issues: read
+ discussions: read
+ packages: read
+ pages: read
+ pull-requests: read
+ repository-projects: read
+ security-events: read
+ statuses: read
jobs:
TagBot:
if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot'
diff --git a/.gitignore b/.gitignore
index b81b1fa0b..a4a7cb627 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+.vscode
*.jl.*.cov
*.jl.cov
*.jl.mem
@@ -5,8 +6,10 @@
/docs/build/
/docs/src/generated/
*output/
-*notebooks/
-examples/test.jl
+scratch/*
+test.jl
+toto.jl
+tata.jl
# Vim swap
*.swp
diff --git a/Project.toml b/Project.toml
index f86099baf..34ee0bc2a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -5,29 +5,50 @@ version = "0.4.1"
[deps]
Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Observables = "510215fc-4207-5dde-b226-833fc4488ee2"
Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
WriteVTK = "64499a7a-5c06-52f2-abe2-ccb03c286192"
+[weakdeps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
+
+[extensions]
+IncompressibleNavierStokesCUDAExt = ["CUDA", "CUDSS"]
+
[compat]
Adapt = "3, 4"
+CairoMakie = "0.11"
+ChainRulesCore = "1"
+CUDA = "5"
+CUDSS = "0.2"
FFTW = "1"
+GLMakie = "0.9"
IterativeSolvers = "0.9"
-Makie = "0.19, 0.20"
+KernelAbstractions = "0.9"
+LinearAlgebra = "1"
+Makie = "0.20"
+NNlib = "0.9"
+Observables = "0.5"
+Printf = "1"
+Random = "1"
+SparseArrays = "1"
+StaticArrays = "1"
Statistics = "1"
WriteVTK = "1"
-julia = "1.7"
+julia = "1.9"
[extras]
-Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Aqua", "CairoMakie", "Test"]
diff --git a/README.md b/README.md
index 1c4e99c40..813106429 100644
--- a/README.md
+++ b/README.md
@@ -9,26 +9,33 @@
This package implements energy-conserving solvers for the incompressible Navier-Stokes
equations on a staggered Cartesian grid. It is based on the Matlab package
-[INS2D](https://github.com/bsanderse/INS2D)/[INS3D](https://github.com/bsanderse/INS3D).
+[INS2D](https://github.com/bsanderse/INS2D)/[INS3D](https://github.com/bsanderse/INS3D). The simulations can be run on the single/multithreaded CPUs or Nvidia GPUs.
+This package also provides experimental support for neural closure models for
+large eddy simulation.
## Installation
To install IncompressibleNavierStokes, open up a Julia-REPL, type `]` to get
into Pkg-mode, and type:
-```
-add IncompressibleNavierStokes
+```julia-repl
+(v1.10) pkg> add IncompressibleNavierStokes
```
which will install the package and all dependencies to your local environment.
-Note that IncompressibleNavierStokes requires Julia version `1.7` or above.
+Note that IncompressibleNavierStokes requires Julia version `1.9` or above.
See the
[Documentation](https://agdestein.github.io/IncompressibleNavierStokes.jl/dev/generated/LidDrivenCavity2D/)
for examples of some typical workflows. More examples can be found in the
[`examples`](examples) directory.
+## Source code for paper
+
+See [here](./lib/PaperDC) for the source code used in the paper
+[Discretize first, filter next: learning divergence-consistent closure models for large-eddy simulation](https://arxiv.org/abs/2403.18088).
+
## Gallery
The velocity and pressure fields may be visualized in a live session using
@@ -37,97 +44,96 @@ The velocity and pressure fields may be visualized in a live session using
snapshot files using the `save_vtk` function, or the full time series using the
`VTKWriter` processor.
-|  |  |  |  |
-|:---------------------------------------:|:-------------------------------------------------------------:|:------------------------------------------------------------:|:-----------------------------------------------------------:|
+|  |  |  |  |
+|:-:|:-:|:-:|:-:|
| [Actuator (2D)](examples/Actuator2D.jl) | [Backward Facing Step (2D)](examples/BackwardFacingStep2D.jl) | [Decaying Turbulence (2D)](examples/DecayingTurbulence2D.jl) | [Taylor-Green Vortex (2D)](examples/TaylorGreenVortex2D.jl) |
-|  |  |  |  |
+|  |  |  |  |
| [Actuator (3D)](examples/Actuator3D.jl) | [Backward Facing Step (3D)](examples/BackwardFacingStep3D.jl) | [Decaying Turbulence (3D)](examples/DecayingTurbulence3D.jl) | [Taylor-Green Vortex (3D)](examples/TaylorGreenVortex3D.jl) |
+IncompressibleNavierStokes also supports adding a temperature equation.
+
+https://github.com/agdestein/IncompressibleNavierStokes.jl/assets/40632532/e4894d97-d890-4d4a-a43a-8eb2772d5130
## Demo
-The following example code using a negative body force on a small rectangle
+The following example code using a negative body force on a small rectangle
with an unsteady inflow. It simulates a wind turbine (actuator) under varying
wind conditions.
+Make sure to have the `GLMakie` and `IncompressibleNavierStokes` installed:
+
+```julia
+using Pkg
+Pkg.add(["GLMakie", "IncompressibleNavierStokes"])
+```
+
+Then run run the following code to make a short animation:
+
```julia
using GLMakie
using IncompressibleNavierStokes
-# Viscosity model
-viscosity_model = LaminarModel(; Re = 100.0)
-
-# Boundary conditions: Unsteady BC requires time derivatives
-u_bc(x, y, t) = x ≈ 0.0 ? cos(π / 6 * sin(π / 6 * t)) : 0.0
-v_bc(x, y, t) = x ≈ 0.0 ? sin(π / 6 * sin(π / 6 * t)) : 0.0
-dudt_bc(x, y, t) = x ≈ 0.0 ? -(π / 6)^2 * cos(π / 6 * t) * sin(π / 6 * sin(π / 6 * t)) : 0.0
-dvdt_bc(x, y, t) = x ≈ 0.0 ? (π / 6)^2 * cos(π / 6 * t) * cos(π / 6 * sin(π / 6 * t)) : 0.0
-bc_type = (;
- u = (; x = (:dirichlet, :pressure), y = (:symmetric, :symmetric)),
- v = (; x = (:dirichlet, :symmetric), y = (:pressure, :pressure)),
+# A 2D grid is a Cartesian product of two vectors
+n = 40
+x = LinRange(0.0, 10.0, 5n + 1)
+y = LinRange(-2.0, 2.0, 2n + 1)
+
+# Boundary conditions
+boundary_conditions = (
+ # Inlet, outlet
+ (
+ # Unsteady BC requires time derivatives
+ DirichletBC(
+ (dim, x, y, t) -> sin(π / 6 * sin(π / 6 * t) + π / 2 * (dim() == 1)),
+ (dim, x, y, t) ->
+ (π / 6)^2 *
+ cos(π / 6 * t) *
+ cos(π / 6 * sin(π / 6 * t) + π / 2 * (dim() == 1)),
+ ),
+ PressureBC(),
+ ),
+
+ # Sides
+ (PressureBC(), PressureBC()),
)
-# A 2D grid is a Cartesian product of two vectors
-x = LinRange(0.0, 10.0, 200)
-y = LinRange(-2.0, 2.0, 80)
-
-# Actuator body force: A thrust coefficient `Cₜ` distributed over a thin rectangle
-xc, yc = 2.0, 0.0 # Disk center
-D = 1.0 # Disk diameter
-δ = 0.11 # Disk thickness
-Cₜ = 5e-4 # Thrust coefficient
-cₜ = Cₜ / (D * δ)
-inside(x, y) = abs(x - xc) ≤ δ / 2 && abs(y - yc) ≤ D / 2
-bodyforce_u(x, y) = -cₜ * inside(x, y)
-bodyforce_v(x, y) = 0.0
+# Actuator body force: A thrust coefficient distributed over a thin rectangle
+inside(x, y) = abs(x - 2.0) ≤ 0.055 && abs(y) ≤ 0.5
+bodyforce(dim, x, y, t) = dim() == 1 && inside(x, y) ? -1.82 : 0.0
# Build setup and assemble operators
-setup = Setup(
- x,
- y;
- viscosity_model,
- u_bc,
- v_bc,
- dudt_bc,
- dvdt_bc,
- bc_type,
- bodyforce_u,
- bodyforce_v,
-);
-
-# Time interval
-t_start, t_end = tlims = (0.0, 40.0)
+setup = Setup(x, y; Re = 100.0, boundary_conditions, bodyforce);
# Initial conditions (extend inflow)
-initial_velocity_u(x, y) = 1.0
-initial_velocity_v(x, y) = 0.0
-initial_pressure(x, y) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t_start;
- initial_pressure,
-);
-
-# Time step processors
-processors = (
- # Record solution every fourth time step
- animator(setup, "vorticity.mp4"; nupdate = 4),
-
- # Log time step information
- step_logger(),
-)
+ustart = create_initial_conditions(setup, (dim, x, y) -> dim() == 1 ? 1.0 : 0.0);
# Solve unsteady Navier-Stokes equations
-V, p, outputs = solve_unsteady(
- setup, V₀, p₀, tlims;
- method = RK44P2(),
- Δt = 0.05,
- processors,
+solve_unsteady(;
+ setup, ustart, tlims = (0.0, 48.0), Δt = 0.05,
+ processors = (
+ anim = animator(; setup, path = "vorticity.mp4", nupdate = 4),
+ log = timelogger(),
+ ),
)
```
The resulting animation is shown below.
https://github.com/agdestein/IncompressibleNavierStokes.jl/assets/40632532/6ee09a03-1674-46e0-843c-000f0b9b9527
+
+## Similar projects
+
+- [WaterLily.jl](https://github.com/weymouth/WaterLily.jl/)
+ Incompressible solver with immersed boundaries
+- [Oceananigans.jl](https://github.com/CliMA/Oceananigans.jl):
+ Ocean simulations
+- [ClimaCore.jl](https://github.com/CliMA/ClimaCore.jl):
+ Atmospheric simulations
+- [Trixi.jl](https://github.com/trixi-framework/Trixi.jl):
+ High order solvers for various hyperbolic equations
+- [Ferrite.jl](https://github.com/Ferrite-FEM/Ferrite.jl):
+ Finite element discretizations
+- [Gridap.jl](https://github.com/gridap/Gridap.jl):
+ Finite element discretizations
+- [FourierFlows.jl](https://github.com/FourierFlows/FourierFlows.jl):
+ Pseudo-spectral discretizations
diff --git a/docs/Project.toml b/docs/Project.toml
index d3322b939..e08de019d 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -1,9 +1,7 @@
[deps]
-CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
-FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
-GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
+Examples = "318dbb63-4243-420f-99f2-d56058123f9d"
IncompressibleNavierStokes = "5e318141-6589-402b-868d-77d7df8c442e"
-LaTeXStrings = "b964fa9f-0449-5b57-a5c2-d3ea65f4040f"
Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
+NeuralClosure = "099dac27-d7f2-4047-93d5-0baee36b9c25"
diff --git a/docs/liveserve.jl b/docs/liveserve.jl
new file mode 100644
index 000000000..0c93059bc
--- /dev/null
+++ b/docs/liveserve.jl
@@ -0,0 +1,9 @@
+push!(LOAD_PATH, "@live-server")
+
+using LiveServer
+
+servedocs(;
+ foldername = @__DIR__,
+ literate = joinpath(@__DIR__, "..", "examples"),
+ skip_dir = joinpath(@__DIR__, "src", "generated"),
+)
diff --git a/docs/make.jl b/docs/make.jl
index 28bf8d07d..b961d031d 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -1,11 +1,21 @@
-# LSP indexing solution
-# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983
-if isdefined(@__MODULE__, :LanguageServer)
- include("../src/IncompressibleNavierStokes.jl")
- using .IncompressibleNavierStokes
-end
+# Show number of threads on GitHub Actions
+@info "" Threads.nthreads()
+
+# Build docs environment
+using Pkg
+Pkg.activate(@__DIR__)
+Pkg.develop([
+ PackageSpec(; path = joinpath(@__DIR__, "..")),
+ PackageSpec(; path = joinpath(@__DIR__, "..", "lib", "NeuralClosure")),
+ PackageSpec(; path = joinpath(@__DIR__, "..", "examples")),
+])
+Pkg.instantiate()
+
+# Get access to example dependencies
+push!(LOAD_PATH, joinpath(@__DIR__, "..", "examples"))
using IncompressibleNavierStokes
+using NeuralClosure
using Literate
using Documenter
using DocumenterCitations
@@ -22,16 +32,16 @@ bib = CitationBibliography(joinpath(@__DIR__, "references.bib"))
# Generate examples
examples = [
"Tutorial: Lid-Driven Cavity (2D)" => "LidDrivenCavity2D",
- "Actuator (2D)" => "Actuator2D",
+ "Convergence: Taylor-Green Vortex (2D)" => "TaylorGreenVortex2D",
+ "Unsteady inflow: Actuator (2D)" => "Actuator2D",
# "Actuator (3D)" => "Actuator3D",
- "Backward Facing Step (2D)" => "BackwardFacingStep2D",
+ "Walls: Backward Facing Step (2D)" => "BackwardFacingStep2D",
# "Backward Facing Step (3D)" => "BackwardFacingStep3D",
"Decaying Turbulunce (2D)" => "DecayingTurbulence2D",
# "Decaying Turbulunce (3D)" => "DecayingTurbulence3D",
# "Lid-Driven Cavity (3D)" => "LidDrivenCavity3D",
"Planar Mixing (2D)" => "PlanarMixing2D",
- "Shear Layer (2D)" => "ShearLayer2D",
- # "Taylor-Green Vortex (2D)" => "TaylorGreenVortex2D",
+ # "Shear Layer (2D)" => "ShearLayer2D",
# "Taylor-Green Vortex (3D)" => "TaylorGreenVortex3D",
]
@@ -45,7 +55,7 @@ for e ∈ examples
end
makedocs(;
- modules = [IncompressibleNavierStokes],
+ modules = [IncompressibleNavierStokes, NeuralClosure],
plugins = [bib],
authors = "Syver Døving Agdestein, Benjamin Sanderse, and contributors",
repo = "https://github.com/agdestein/IncompressibleNavierStokes.jl/blob/{commit}{path}#{line}",
@@ -66,12 +76,14 @@ makedocs(;
"Time discretization" => "equations/time.md",
],
"Features" => [
+ "Boundary conditions" => "features/bc.md",
+ "Pressure solvers" => "features/pressure.md",
"Floating point precision" => "features/precision.md",
"GPU Support" => "features/gpu.md",
+ "Operators" => "features/operators.md",
+ "Temperature equation" => "features/temperature.md",
"Large eddy simulation" => "features/les.md",
- "Pressure solvers" => "features/pressure.md",
- "Boundary conditions" => "features/bc.md",
- "Time steppers" => "features/steppers.md",
+ "Neural closure models" => "features/closure.md",
],
"API Reference" =>
["API" => "api/api.md", "Runge-Kutta methods" => "api/tableaux.md"],
diff --git a/docs/references.bib b/docs/references.bib
index 513805881..2a45b8834 100644
--- a/docs/references.bib
+++ b/docs/references.bib
@@ -1,3 +1,27 @@
+@article{San2012,
+ title={High-order methods for decaying two-dimensional homogeneous isotropic turbulence},
+ volume={63},
+ ISSN={0045-7930},
+ url={http://dx.doi.org/10.1016/j.compfluid.2012.04.006},
+ DOI={10.1016/j.compfluid.2012.04.006},
+ journal={Computers \& Fluids},
+ publisher={Elsevier BV},
+ author={San, Omer and Staples, Anne E.},
+ year={2012},
+ month={jun},
+ pages={105–127}
+}
+@misc{Sanderse2023,
+ Author = {Benjamin Sanderse and Francesc Xavier Trias},
+ Title = {Energy-consistent discretization of viscous dissipation with application
+ to natural convection flow},
+ Eprint = {2307.10874v1},
+ ArchivePrefix = {arXiv},
+ PrimaryClass = {physics.flu-dyn},
+ Year = {2023},
+ Month = {Jul},
+ Url = {http://arxiv.org/abs/2307.10874v1},
+}
@article{Beck2021,
author = {Beck, Andrea and Kurz, Marius},
doi = {10.1002/gamm.202100002},
@@ -50,6 +74,18 @@ @article{Harlow1965
URL = {https://aip.scitation.org/doi/abs/10.1063/1.1761178},
eprint = {https://aip.scitation.org/doi/pdf/10.1063/1.1761178},
}
+@article{Jeong1995,
+ doi = {10.1017/s0022112095000462},
+ url = {https://doi.org/10.1017%2Fs0022112095000462},
+ year = {1995},
+ month = {feb},
+ publisher = {Cambridge University Press ({CUP})},
+ volume = {285},
+ pages = {69--94},
+ author = {Jinhee Jeong and Fazle Hussain},
+ title = {On the identification of a vortex},
+ journal = {J. Fluid. Mech.}
+}
@article{Kochkov2021,
author = {Dmitrii Kochkov and Jamie A. Smith and Ayya Alieva and Qing Wang and Michael P. Brenner and Stephan Hoyer},
doi = {10.1073/pnas.2101784118},
@@ -68,10 +104,22 @@ @article{Kurz2022
author = {Kurz, Marius and Offenhäuser, Philipp and Beck, Andrea},
keywords = {Fluid Dynamics (physics.flu-dyn), Artificial Intelligence (cs.AI), Computational Engineering, Finance, and Science (cs.CE), Machine Learning (cs.LG), FOS: Physical sciences, FOS: Physical sciences, FOS: Computer and information sciences, FOS: Computer and information sciences},
title = {Deep Reinforcement Learning for Turbulence Modeling in Large Eddy Simulations},
- publisher = {arXiv},
+ publisher = {arXiv:2206.11038},
+ journal = {arXiv},
year = {2022},
copyright = {Creative Commons Attribution Non Commercial No Derivatives 4.0 International}
}
+@article{LiJiajia2019,
+ Author = {Jiajia Li and Pablo M. Carrica},
+ Title = {A simple approach for vortex core visualization},
+ journal = {arXiv:1910.06998},
+ Eprint = {1910.06998},
+ ArchivePrefix = {arXiv},
+ PrimaryClass = {physics.flu-dyn},
+ Year = {2019},
+ Month = {Oct},
+ Url = {http://arxiv.org/abs/1910.06998},
+}
@article{List2022,
author = {List, Björn and Chen, Li-Wei and Thuerey, Nils},
copyright = {arXiv.org perpetual, non-exclusive license},
@@ -92,6 +140,13 @@ @misc{MacArt2021
url = {https://arxiv.org/abs/2105.01030},
year = {2021},
}
+@book{Orlandi2000,
+ title={Fluid flow phenomena: a numerical toolkit},
+ author={Orlandi, Paolo},
+ volume={55},
+ year={2000},
+ publisher={Springer Science \& Business Media}
+}
@book{Sagaut2006,
address = {Berlin/Heidelberg},
author = {Sagaut, Pierre},
@@ -105,7 +160,7 @@ @book{Sagaut2006
}
@article{Sanderse2012,
author = {B. Sanderse and B. Koren},
- doi = {https://doi.org/10.1016/j.jcp.2011.11.028},
+ doi = {10.1016/j.jcp.2011.11.028},
issn = {0021-9991},
journal = {Journal of Computational Physics},
keywords = {Differential–algebraic equations, Incompressible Navier–Stokes equations, Temporal accuracy, Time integration, Runge–Kutta method, Moving meshes},
@@ -118,7 +173,7 @@ @article{Sanderse2012
}
@article{Sanderse2013,
author = {B. Sanderse},
- doi = {https://doi.org/10.1016/j.jcp.2012.07.039},
+ doi = {10.1016/j.jcp.2012.07.039},
issn = {0021-9991},
journal = {Journal of Computational Physics},
keywords = {Energy conservation, Time reversibility, Runge–Kutta method, Additive Runge–Kutta method, Incompressible Navier–Stokes equations, Differential–algebraic equations, Algebraic stability, stability, Stiffness},
@@ -130,7 +185,7 @@ @article{Sanderse2013
}
@article{Sanderse2014,
author = {B. Sanderse and R.W.C.P. Verstappen and B. Koren},
- doi = {https://doi.org/10.1016/j.jcp.2013.10.002},
+ doi = {10.1016/j.jcp.2013.10.002},
issn = {0021-9991},
journal = {Journal of Computational Physics},
keywords = {Incompressible Navier–Stokes equations, Symmetry preservation, Energy conservation, Boundary conditions, Fourth order accuracy},
@@ -153,6 +208,19 @@ @article{Sanderse2020
volume = {421},
year = {2020}
}
+@article{Silvis2017,
+ title = {Physical consistency of subgrid-scale models for large-eddy simulation of incompressible turbulent flows},
+ volume = {29},
+ ISSN = {1089-7666},
+ url = {http://dx.doi.org/10.1063/1.4974093},
+ DOI = {10.1063/1.4974093},
+ number = {1},
+ journal = {Physics of Fluids},
+ publisher = {AIP Publishing},
+ author = {Silvis, Maurits H. and Remmerswaal, Ronald A. and Verstappen, Roel},
+ year = {2017},
+ month = jan,
+}
@article{Verstappen1997,
doi = {10.1023/a:1004255329158},
url = {https://doi.org/10.1023%2Fa%3A1004255329158},
diff --git a/docs/src/api/api.md b/docs/src/api/api.md
index 9e71f2665..1f51ef873 100644
--- a/docs/src/api/api.md
+++ b/docs/src/api/api.md
@@ -9,18 +9,6 @@ IncompressibleNavierStokes
Setup
```
-## Boundary conditions
-
-```@docs
-BoundaryConditions
-get_bc_vectors
-```
-
-## Force
-
-```@docs
-SteadyBodyForce
-```
## Grid
@@ -32,78 +20,6 @@ max_size
stretched_grid
```
-## Visocosity Models
-
-```@docs
-AbstractViscosityModel
-LaminarModel
-MixingLengthModel
-SmagorinskyModel
-QRModel
-```
-
-## Convection Models
-
-```@docs
-AbstractConvectionModel
-NoRegConvectionModel
-C2ConvectionModel
-C4ConvectionModel
-LerayConvectionModel
-```
-
-## Momentum
-
-```@docs
-MomentumCache
-check_symmetry
-compute_conservation
-convection
-convection!
-convection_components
-convection_components!
-diffusion
-diffusion!
-momentum
-momentum!
-momentum_allstage
-momentum_allstage!
-strain_tensor
-turbulent_K
-turbulent_viscosity
-```
-
-## Operators
-
-```@docs
-Operators
-operator_averaging
-operator_convection_diffusion
-operator_divergence
-operator_interpolation
-operator_postprocessing
-operator_regularization
-operator_turbulent_diffusion
-operator_viscosity
-operator_filter
-```
-
-## Postprocess
-
-```@docs
-get_streamfunction
-get_velocity
-get_vorticity
-vorticity!
-plot_force
-plot_grid
-plot_pressure
-plot_streamfunction
-plot_velocity
-plot_vorticity
-save_vtk
-```
-
## Preprocess
```@docs
@@ -114,12 +30,14 @@ random_field
## Processors
```@docs
-step_logger
+processor
+timelogger
vtk_writer
-field_saver
-field_plotter
-energy_history_plotter
-energy_spectrum_plotter
+fieldsaver
+realtimeplotter
+fieldplot
+energy_history_plot
+energy_spectrum_plot
animator
```
@@ -131,52 +49,11 @@ solve_unsteady
solve_steady_state
```
-### Pressure solvers
-
-```@docs
-AbstractPressureSolver
-DirectPressureSolver
-CGPressureSolver
-SpectralPressureSolver
-pressure_additional_solve
-pressure_additional_solve!
-pressure_poisson
-pressure_poisson!
-```
-
-## Time steppers
-
-```@docs
-AbstractODEMethod
-AbstractRungeKuttaMethod
-AdamsBashforthCrankNicolsonMethod
-OneLegMethod
-ExplicitRungeKuttaMethod
-ImplicitRungeKuttaMethod
-
-isexplicit
-lambda_conv_max
-lambda_diff_max
-nstage
-ode_method_cache
-runge_kutta_method
-step
-step!
-```
-
-## Filter
-```@docs
-create_top_hat_u
-create_top_hat_v
-create_top_hat_p
-create_top_hat_velocity
-```
-
## Utils
```@docs
-filter_convection
-filter_convection!
+save_vtk
+plotgrid
get_lims
plotmat
```
diff --git a/docs/src/api/tableaux.md b/docs/src/api/tableaux.md
index f55b15640..9c0c29b9f 100644
--- a/docs/src/api/tableaux.md
+++ b/docs/src/api/tableaux.md
@@ -4,93 +4,96 @@ CurrentModule = IncompressibleNavierStokes
# API Reference -- Runge-Kutta methods
+```@docs
+RKMethods
+```
## Explicit Methods
```@docs
-FE11
-SSP22
-SSP42
-SSP33
-SSP43
-SSP104
-rSSPs2
-rSSPs3
-Wray3
-RK56
-DOPRI6
+RKMethods.FE11
+RKMethods.SSP22
+RKMethods.SSP42
+RKMethods.SSP33
+RKMethods.SSP43
+RKMethods.SSP104
+RKMethods.rSSPs2
+RKMethods.rSSPs3
+RKMethods.Wray3
+RKMethods.RK56
+RKMethods.DOPRI6
```
## Implicit Methods
```@docs
-BE11
-SDIRK34
-ISSPm2
-ISSPs3
+RKMethods.BE11
+RKMethods.SDIRK34
+RKMethods.ISSPm2
+RKMethods.ISSPs3
```
## Half explicit methods
```@docs
-HEM3
-HEM3BS
-HEM5
+RKMethods.HEM3
+RKMethods.HEM3BS
+RKMethods.HEM5
```
## Classical Methods
```@docs
-GL1
-GL2
-GL3
-RIA1
-RIA2
-RIA3
-RIIA1
-RIIA2
-RIIA3
-LIIIA2
-LIIIA3
+RKMethods.GL1
+RKMethods.GL2
+RKMethods.GL3
+RKMethods.RIA1
+RKMethods.RIA2
+RKMethods.RIA3
+RKMethods.RIIA1
+RKMethods.RIIA2
+RKMethods.RIIA3
+RKMethods.LIIIA2
+RKMethods.LIIIA3
```
## Chebyshev methods
```@docs
-CHDIRK3
-CHCONS3
-CHC3
-CHC5
+RKMethods.CHDIRK3
+RKMethods.CHCONS3
+RKMethods.CHC3
+RKMethods.CHC5
```
## Miscellaneous Methods
```@docs
-Mid22
-MTE22
-CN22
-Heun33
-RK33C2
-RK33P2
-RK44
-RK44C2
-RK44C23
-RK44P2
+RKMethods.Mid22
+RKMethods.MTE22
+RKMethods.CN22
+RKMethods.Heun33
+RKMethods.RK33C2
+RKMethods.RK33P2
+RKMethods.RK44
+RKMethods.RK44C2
+RKMethods.RK44C23
+RKMethods.RK44P2
```
## DSRK Methods
```@docs
-DSso2
-DSRK2
-DSRK3
+RKMethods.DSso2
+RKMethods.DSRK2
+RKMethods.DSRK3
```
## "Non-SSP" Methods of Wong & Spiteri
```@docs
-NSSP21
-NSSP32
-NSSP33
-NSSP53
+RKMethods.NSSP21
+RKMethods.NSSP32
+RKMethods.NSSP33
+RKMethods.NSSP53
```
diff --git a/docs/src/assets/grid.png b/docs/src/assets/grid.png
index ed11cd750..a40dbc865 100644
Binary files a/docs/src/assets/grid.png and b/docs/src/assets/grid.png differ
diff --git a/docs/src/assets/grid.svg b/docs/src/assets/grid.svg
deleted file mode 100644
index 9e4be4ce2..000000000
--- a/docs/src/assets/grid.svg
+++ /dev/null
@@ -1,1126 +0,0 @@
-
-
-
-
diff --git a/docs/src/equations/ns.md b/docs/src/equations/ns.md
index 769f96646..326dd13e4 100644
--- a/docs/src/equations/ns.md
+++ b/docs/src/equations/ns.md
@@ -1,92 +1,61 @@
# Incompressible Navier-Stokes equations
-Let ``d \in \{2, 3\}`` denote the spatial dimension, and ``\Omega \subset
-\mathbb{R}^d`` some spatial domain. The incompressible Navier-Stokes equations
-on ``\Omega`` are comprised of a mass equation and ``d`` momentum equations. In
-conservative form, they are given by
+The incompressible Navier-Stokes equations describe conservation of mass and
+conservation of momentum, which can be written as a divergence-free constraint
+and an evolution equation:
```math
-\begin{align*}
-\nabla \cdot u & = 0, \\
-\frac{\partial u}{\partial t} + \nabla \cdot (u u^\mathsf{T}) & = -\nabla p +
-\nu \nabla^2 u + f.
-\end{align*}
+\begin{aligned}
+ \nabla \cdot u & = 0, \\
+ \frac{\partial u}{\partial t} + \nabla \cdot (u u^\mathsf{T})
+ & = -\nabla p + \nu \nabla^2 u + f,
+\end{aligned}
```
-where ``u = (u^1, \dots, u^d)`` is the velocity field, ``p`` is the
+where ``\Omega \subset \mathbb{R}^d`` is the domain, ``d \in \{2, 3\}`` is the
+spatial dimension, ``u = (u^1, \dots, u^d)`` is the velocity field, ``p`` is the
pressure, ``\nu`` is the kinematic viscosity, and ``f = (f^1, \dots, f^d)`` is
-the body force per unit of volume. In scalar notation, the equations can be
-written as
-
-```math
-\begin{align*}
-\sum_{\beta = 1}^d \frac{\partial u^\beta}{\partial x^\beta} & = 0, \\
-\frac{\partial u^\alpha}{\partial t} + \sum_{\beta = 1}^d
-\frac{\partial}{\partial x^\beta} (u^\alpha u^\beta) & = -\frac{\partial
-p}{\partial x^\alpha} + \nu \sum_{\beta = 1}^d \frac{\partial^2 u^\alpha}{\partial
-x^\beta \partial x^\beta} + f^\alpha.
-\end{align*}
-```
-
+the body force per unit of volume. The velocity, pressure, and body force are
+functions of the spatial coordinate ``x = (x^1, \dots, x^d)`` and time ``t``.
+We assume that ``\Omega`` is a rectangular domain.
## Integral form
-When discretizing the Navier-Stokes equations it can be useful to integrate the
-equations over an arbitrary control volume ``\mathcal{O}``. Its boundary will
-be denoted ``\partial \mathcal{O}``, with normal ``n`` and surface element
-``\mathrm{d} \Gamma``.
-
-The mass equation in integral form is given by
-
-```math
-\int_{\partial \mathcal{O}} u \cdot n \, \mathrm{d} \Gamma = 0,
-```
-
-where we have used the divergence theorem to convert the volume integral to a
-surface integral. Similarly, the momentum equations take the form
+The integral form of the Navier-Stokes equations is used as starting point to
+develop a spatial discretization:
```math
-\frac{\partial }{\partial t} \int_\mathcal{O} u \, \mathrm{d} \Omega
-= \int_{\partial \mathcal{O}} \left( - u u^\mathsf{T} - P + \nu S \right) \cdot n \,
-\mathrm{d} \Gamma + \int_\mathcal{O} f \mathrm{d} \Omega
+\begin{aligned}
+ \frac{1}{|\mathcal{O}|} \int_{\partial \mathcal{O}}
+ u \cdot n \, \mathrm{d} \Gamma & = 0, \\
+ \frac{\mathrm{d} }{\mathrm{d} t} \frac{1}{|\mathcal{O}|}
+ \int_\mathcal{O} u \, \mathrm{d} \Omega
+ & = \frac{1}{|\mathcal{O}|} \int_{\partial \mathcal{O}}
+ \left( - u u^\mathsf{T} - p I + \nu \nabla u \right) \cdot n \,
+ \mathrm{d} \Gamma +
+ \frac{1}{|\mathcal{O}|}\int_\mathcal{O} f \mathrm{d} \Omega,
+\end{aligned}
```
-where ``P = p \mathrm{I}`` is the hydrostatic stress tensor
-and ``S = \nabla u + (\nabla u)^\mathsf{T}`` is the strain-rate tensor. Here we
-have once again used the divergence theorem.
-
-!!! note "Strain-rate tensor"
- The second term in the strain rate tensor is optional, as
-
- ```math
- \int_{\partial \mathcal{O}} (\nabla u)^\mathsf{T} \cdot n \, \mathrm{d} \Gamma = 0
- ```
-
- due to the divergence freeness of ``u`` (mass equation).
-
+where ``\mathcal{O} \subset \Omega`` is an arbitrary control volume with boundary
+``\partial \mathcal{O}``, normal ``n``, surface element ``\mathrm{d} \Gamma``, and
+volume size ``|\mathcal{O}|``. We have divided by the control volume sizes in the
+integral form.
## Boundary conditions
-Consider a part ``\Gamma`` of the total boundary ``\partial \Omega``, with
-normal ``n``. We allow for the following types of boundary conditions on
-``\Gamma``:
-
-- Periodic boundary conditions: ``u(x) = u(x + \tau)`` and ``p(x) = p(x + \tau)`` for ``x \in
- \Gamma``, where ``\tau`` is a constant translation to another part of the
- boundary ``\partial \Omega``. This usually requires ``\Gamma`` and its
- periodic counterpart ``\Gamma + \tau`` to be flat and rectangular (including
- in this software suite).
-- Dirichlet boundary conditions: ``u = u_\text{BC}`` on ``\Gamma``. A common
- situation here is that of a sticky wall, with "no slip boundary conditions.
- Then ``u_\text{BC} = 0``.
-- Symmetric boundary conditions: Same velocity field at each side. This implies
- zero Dirichlet conditions for the normal component (``n \cdot u = 0``), and
- zero Neumann conditions for the parallel component: ``n \cdot \nabla (u - (n
- \cdot u) n) = 0``.
-- Pressure boundary conditions: The pressure is prescribed (``p =
- p_\text{BC}``), while the velocity has zero Neumann conditions:
- ``n \cdot \nabla u = 0``.
-
+The boundary conditions on a part of the boundary
+``\Gamma \subset \partial \Omega`` are one or more of the following:
+
+- Dirichlet: ``u = u_\text{BC}`` on ``\Gamma`` for some ``u_\text{BC}``;
+- Neumann: ``\nabla u \cdot n = 0`` on ``\Gamma``;
+- Periodic: ``u(x) = u(x + \tau)`` and ``p(x) = p(x + \tau)``
+ for ``x \in \Gamma``, where
+ ``\Gamma + \tau \subset \partial \Omega``
+ is another part of the boundary and
+ ``\tau`` is a translation vector;
+- Stress free: ``\sigma \cdot n = 0`` on ``\Gamma``,
+ where ``\sigma = \left(- p \mathrm{I} + 2 \nu S \right)``.
## Pressure equation
diff --git a/docs/src/equations/spatial.md b/docs/src/equations/spatial.md
index ee274bfd6..e98183f68 100644
--- a/docs/src/equations/spatial.md
+++ b/docs/src/equations/spatial.md
@@ -1,391 +1,173 @@
# Spatial Discretization
-To discretize the incompressible Navier-Stokes equations, we will use finite
-volumes on a staggered Cartesian grid, as proposed by Harlow and Welsh
-[Harlow1965](@cite). We will use the notation of Sanderse [Sanderse2012](@cite)
-[Sanderse2013](@cite) [Sanderse2014](@cite).
-
-Let ``d \in \{2, 3\}`` denote the spatial dimension (2D or 3D). We will make
-use of the "Cartesian" index ``I = (i, j)`` in 2D or ``I = (i, j, k)`` in 3D,
-with ``I(1) = i``, ``I(2) = j``, and ``I(3) = k``. Here, the indices ``I``,
-``i``, ``j``, and ``k``, represent discrete degrees of freedom, and can take
-integer or half values (e.g. ``3`` or ``5/2``). To specify a
-spatial dimension, we will use the symbols ``(\alpha, \beta, \gamma) \in \{1,
-\dots, d\}^3``. We will use the symbol ``\delta(\alpha) = (\delta_{\alpha
-\beta})_{\beta = 1}^d \in \{0, 1\}^d`` to indicate a perturbation in the
-direction ``\alpha``, where ``\delta_{\alpha \beta}`` is the Kronecker symbol.
-The spatial variable is ``x = (x^1, \dots, x^d) \in \Omega \subset
-\mathbb{R}^d``.
-
-
-## Finite volumes
-
-We here assume that ``\Omega = \prod_{\alpha = 1}^d [0, L^\alpha]`` has the
-shape of a box with side lengths ``L^\alpha > 0``. This allows for partitioning
-``\Omega`` into the finite volumes
+the ``d`` spatial dimensions are indexed by ``\alpha \in \{1, \dots, d\}``. The
+``\alpha``-th unit vector is denoted ``e_\alpha = (e_{\alpha \beta})_{\beta =
+1}^d``, where the Kronecker symbol ``e_{\alpha \beta}`` is ``1`` if ``\alpha =
+\beta`` and ``0`` otherwise. We note ``h_\alpha = e_\alpha / 2``. The Cartesian
+index ``I = (I_1, \dots, I_d)`` is used to avoid repeating terms and equations
+``d`` times, where ``I_\alpha`` is a scalar index (typically one of ``i``,
+``j``, and ``k`` in common notation). This notation is dimension-agnostic, since
+we can write ``u_I`` instead of ``u_{i j}`` in 2D or ``u_{i j k}`` in 3D. In our
+Julia implementation of the solver we use the same Cartesian notation
+(`u[I]` instead of `u[i, j]` or `u[i, j, k]`).
+
+For the discretization scheme, we use a staggered Cartesian grid as proposed by
+Harlow and Welch [Harlow1965](@cite). Consider a rectangular domain
+``\Omega = \prod_{\alpha = 1}^d [a_\alpha, b_\alpha]``, where ``a_\alpha <
+b_\alpha`` are the domain boundaries and ``\prod`` is a Cartesian product. Let
+``\Omega = \bigcup_{I \in \mathcal{I}} \Omega_I`` be a partitioning of ``\Omega``,
+where ``\mathcal{I} = \prod_{\alpha = 1}^d \{ \frac{1}{2}, 2 - \frac{1}{2},
+\dots, N_\alpha - \frac{1}{2} \}`` are volume center indices, ``N = (N_1, \dots,
+N_d) \in \mathbb{N}^d`` are the number of volumes in each dimension,
+``\Omega_I = \prod_{\alpha = 1}^d \Delta^\alpha_{I_\alpha}`` is a finite
+volume, ``\Gamma^\alpha_I = \Omega_{I - h_\alpha} \cap \Omega_{I + h_\alpha} =
+\prod_{\beta \neq \alpha} \Delta^\beta_{I_\beta}`` is a volume face,
+``\Delta^\alpha_i = \left[ x^\alpha_{i - \frac{1}{2}}, x^\alpha_{i + \frac{1}{2}}
+\right]`` is a volume edge, ``x^\alpha_0, \dots, x^\alpha_{N_\alpha}`` are volume
+boundary coordinates, and ``x^\alpha_i = \frac{1}{2} \left(x^\alpha_{i -
+\frac{1}{2}} + x^\alpha_{i + \frac{1}{2}}\right)`` for ``i \in \{ 1 / 2, \dots,
+N_\alpha - 1 / 2\}`` are volume center coordinates. We also define the operator
+``\delta_\alpha`` which maps a discrete scalar field ``\varphi = (\varphi_I)_I`` to
```math
-\Omega_I = \prod_{\alpha = 1}^d \left[ x^\alpha_{I(\alpha) - \frac{1}{2}},
-x^\alpha_{I(\alpha) + \frac{1}{2}} \right], \quad I \in \mathcal{I}.
+(\delta_\alpha \varphi)_I = \frac{\varphi_{I + h_\alpha} - \varphi_{I -
+h_\alpha}}{| \Delta^\alpha_{I_\alpha} |}.
```
-Just like ``\Omega`` itself, they represent rectangles in 2D and prisms in 3D.
-They are fully defined by the vectors of volume face coordinates ``x^\alpha =
-\left( x^\alpha_{i} \right)_{i = 0}^{N(\alpha)} \in \mathbb{R}^{N(\alpha) +
-1}``, where ``N = (N(1), \dots, N(d)) \in \mathbb{N}^d`` are the numbers of
-finite volumes in each dimension and ``\mathcal{I} = \prod_{\alpha = 1}^d
-\left\{ \frac{1}{2}, 2 - \frac{1}{2}, \dots, N(\alpha) - \frac{1}{2} \right\}``
-the set of finite volume indices (note that the reference volumes are indexed
-by half indices only). The components ``x^\alpha_{i}`` are not assumed to be
-uniformly spaced. But we do assume that they are strictly increasing with
-``i``, with ``x^\alpha_0 = 0`` and ``x^\alpha_{N(\alpha)} = L^\alpha``.
-
-The coordinates of the volume centers are determined from the those of the
-volume boundaries by
-``x^\alpha_{I(\alpha)} = \frac{1}{2} (x^\alpha_{I(\alpha) - \frac{1}{2}} +
-x_{I(\alpha) + \frac{1}{2}})``
-for ``I \in \mathcal{I}``. This allows for defining shifted volumes such as
-``\Omega_{I + \delta(\alpha) / 2}`` and
-``\Omega_{I + \delta(\alpha) / 2 + \delta(\beta) / 2}``. The original volumes
-(with indices in ``\mathcal{I}``) will be referred to as the reference finite
-volumes.
-
-We also define the volume widths/depths/heights ``\Delta x^\alpha_i =
-x^\alpha_{i + \frac{1}{2}} - x^\alpha_{i - \frac{1}{2}}``, where ``i`` can take
-half values. The volume sizes are thus ``| \Omega_{I} | = \prod_{\alpha = 1}^d
-\Delta x^\alpha_{I(\alpha)}``.
-
-In addition to the finite volumes and their shifted variants, we
-define the surface
-
-```math
-\Gamma^\alpha_I = \prod_{\beta = 1}^d \begin{cases}
- \left\{ x^\beta_{I(\beta)} \right\}, & \quad \alpha = \beta \\
- \left[ x^\beta_{I(\beta) - 1 / 2}, x^\beta_{I(\beta) + 1 / 2} \right], & \quad
- \text{otherwise},
-\end{cases}
-```
-where ``I`` can take integer or half-values. It is the interface between
-``\Omega_{I - \delta(\alpha) / 2}`` and ``\Omega_{I + \delta(\alpha) / 2}``,
-and has surface normal ``\delta(\alpha)``.
-
-In each reference finite volume ``\Omega_{I}`` (``I \in \mathcal{I}``), there
-are three different types of positions in which quantities of interest can be
-defined:
-
-- The volume center ``x_I = (x_{I(1)}, \dots, x_{I(d)})``, where the discrete
- pressure ``p_I`` is defined;
-- The right/rear/top volume face centers ``x_{I + \delta(\alpha) / 2}``, where
- the discrete ``\alpha``-velocity component ``u^\alpha_{I + \delta(\alpha) / 2}`` is defined;
-- The right-rear-top volume corner ``x_{I + \sum_{\alpha} \delta(\alpha) /
- 2}``, where the discrete vorticity ``\omega_{I + \sum_{\alpha} \delta(\alpha) /
- 2}`` is defined.
-
-The vectors of unknowns ``u^\alpha_h`` and ``p_h`` will not contain all the
-half-index components, only those from their own canonical position. The
-unknown discrete pressure represents the average pressure in each reference
-volume, and the unknown discrete velocity components represent exchange of mass
-between neighboring volumes.
-
-In 2D, this finite volume configuration is illustrated as follows:
+It can be interpreted as a discrete equivalent of the continuous operator
+``\frac{\partial}{\partial x^\alpha}``. All the above definitions are extended to
+be valid in volume centers ``I \in \mathcal{I}``, volume faces
+``I \in \mathcal{I} + h_\alpha``,
+or volume corners ``I \in \mathcal{I} + \sum_{\alpha = 1}^d h_\alpha``.
+The discretization is illustrated below:

-
-## Interpolation
-
-When a quantity is required *outside* of its native point, we will use interpolation. Examples:
-
-- To compute ``u^\alpha`` at the pressure point ``x_I``, ``I \in \mathcal{I}``:
- ```math
- \begin{split}
- u^\alpha_I & = \frac{x^\alpha_{I(\alpha) + 1 / 2} -
- x^\alpha_{I(\alpha)}}{x^\alpha_{I(\alpha) + 1 / 2} - x^\alpha_{I(\alpha) - 1 / 2}}
- u_{I - \delta(\alpha) / 2}
- + \frac{x^\alpha_{I(\alpha)} - x^\alpha_{I(\alpha) - 1 / 2}}{x^\alpha_{I(\alpha) + 1 / 2} - x^\alpha_{I(\alpha) - 1 / 2}}
- u_{I + \delta(\alpha) / 2} \\
- & = \frac{1}{2} \left( u_{I - \delta(\alpha) / 2} + u_{I + \delta(\alpha) / 2} \right).
- \end{split}
- ```
- Interpolation weights from volume faces to volume centers are always
- ``\frac{1}{2}``.
-- To compute ``u^\alpha`` at center of edge between ``\alpha``-face and
- ``\beta``-face ``x_{I + \delta(\alpha) / 2 + \delta(\beta) / 2}``:
- ```math
- u^\alpha_{I + \delta(\alpha) / 2 + \delta(\beta) / 2} =
- \frac{x^\beta_{I(\beta) + 1} - x^\beta_{I(\beta) + 1 / 2}}{x^\beta_{I(\beta) + 1} - x^\beta_{I(\beta)}}
- u^\alpha_{I + \delta(\alpha) / 2}
- + \frac{x^\beta_{I(\beta) + 1 / 2} - x^\beta_{I(\beta)}}{x^\beta_{I(\beta) + 1} - x^\beta_{I(\beta)}}
- u^\alpha_{I + \delta(\alpha) / 2 + \delta(\beta)}.
- ```
- Note that the grid is allowed to be non-uniform, so the interpolation weights
- may unequal and different from ``\frac{1}{2}``.
-- To compute ``p`` at ``u^\alpha``-points:
- ```math
- p_{I + \delta(\alpha) / 2} =
- \frac{x^\alpha_{I(\alpha) + 1} - x^\alpha_{I(\alpha) + 1 / 2}}{x^\alpha_{I(\alpha) + 1} - x^\alpha_{I(\alpha)}}
- p_{I}
- + \frac{x^\alpha_{I(\alpha) + 1 / 2} - x^\alpha_{I(\alpha)}}{x^\alpha_{I(\alpha) + 1} - x^\alpha_{I(\alpha)}}
- p_{I + \delta(\alpha)}
- ```
-
-
## Finite volume discretization of the Navier-Stokes equations
-We will consider the integral form of the Navier-Stokes equations. This has the
-advantage that some of the spatial derivatives dissapear, reducing the amount
-of finite difference approximations we need to perform.
-
-
-### Mass equation
+We now define the unknown degrees of freedom. The average pressure in
+``\Omega_I``, ``I \in \mathcal{I}`` is approximated by the quantity ``p_I(t)``. The
+average ``\alpha``-velocity on the face ``\Gamma^\alpha_I``, ``I \in \mathcal{I} +
+h_\alpha`` is approximated by the quantity ``u^\alpha_I(t)``. Note how the pressure
+``p`` and the ``d`` velocity fields ``u^\alpha`` are each defined in their own
+canonical positions ``x_I`` and ``x_{I + h_\alpha}`` for ``I \in \mathcal{I}``.
+In the following, we derive equations for these unknowns.
-The mass equation takes the form
+Using the pressure control volume ``\mathcal{O} = \Omega_I`` with ``I \in
+\mathcal{I}`` in the mass integral constraint and
+approximating the face integrals with the mid-point quadrature rule
+``\int_{\Gamma_I} u \, \mathrm{d} \Gamma \approx | \Gamma_I | u_I`` results in the
+discrete divergence-free constraint
```math
-\int_{\partial \mathcal{O}} u \cdot n \, \mathrm{d} \Gamma = 0, \quad \forall
-\mathcal{O} \subset \Omega.
+\sum_{\alpha = 1}^d (\delta_\alpha u^\alpha)_I = 0.
```
-Using the pressure volume ``\mathcal{O} = \Omega_{I}``, we get
+Note how dividing by the volume size results in a discrete equation resembling
+the continuous one (since ``| \Omega_I | = | \Gamma^\alpha_I | |
+\Delta^\alpha_{I_\alpha} |``).
-```math
-\sum_{\alpha = 1}^d \left( \int_{\Gamma^\alpha_{I + \delta(\alpha) / 2}}
-u^\alpha \, \mathrm{d} \Gamma - \int_{\Gamma_{I - \delta(\alpha) / 2}^\alpha}
-u^\alpha \, \mathrm{d} \Gamma
-\right) = 0.
-```
-
-Assuming that the flow is fully resolved, meaning that ``\Omega_{I}`` is is
-sufficiently small such that ``u`` is locally linear, we can perform the
-local approximation (quadrature)
+Similarly, choosing an ``\alpha``-velocity control volume ``\mathcal{O} =
+\Omega_{I}`` with ``I \in \mathcal{I} + h_\alpha`` in the integral momentum
+equation, approximating the volume- and face integrals using
+the mid-point quadrature rule, and replacing remaining spatial derivatives in
+the diffusive term with a finite difference approximation gives the discrete
+momentum equations
```math
-\int_{\Gamma^\alpha_I} u^\alpha \, \mathrm{d} \Gamma \approx | \Gamma^\alpha_I
-| u^\alpha_{I}.
+\frac{\mathrm{d}}{\mathrm{d} t} u^\alpha_{I} =
+- \sum_{\beta = 1}^d
+(\delta_\beta (u^\alpha u^\beta))_{I}
++ \nu \sum_{\beta = 1}^d
+(\delta_\beta \delta_\beta u^\alpha)_{I}
++ f^\alpha(x_{I})
+- (\delta_\alpha p)_{I}.
```
-This yields the discrete mass equation
+where we made the assumption that ``f`` is constant in time for simplicity.
+The outer discrete derivative in ``(\delta_\beta \delta_\beta u^\alpha)_{I}`` is
+required at the position ``I``, which means that the inner derivative is evaluated
+as ``(\delta_\beta u^\alpha)_{I + h_\beta}`` and ``(\delta_\beta u^\alpha)_{I -
+h_\beta}``, thus requiring ``u^\alpha_{I - 2 h_\beta}``, ``u^\alpha_{I}``, and
+``u^\alpha_{I + 2 h_\beta}``, which are all in their canonical positions. The two
+velocity components in the convective term ``u^\alpha u^\beta`` are required at
+the positions ``I - h_\beta`` and ``I + h_\beta``, which are outside the canonical
+positions. Their value at the required position is obtained using averaging with
+weights ``1 / 2`` for the ``\alpha``-component and with linear interpolation for the
+``\beta``-component. This preserves the skew-symmetry of the convection operator,
+such that energy is conserved (in the convective term) [Verstappen2003](@cite).
-```math
-\sum_{\alpha = 1}^d
-\left| \Gamma^\alpha_I \right| \left( u^\alpha_{I + \delta(\alpha) / 2} -
-u^\alpha_{I - \delta(\alpha) / 2} \right) = 0.
-```
+## Boundary conditions
-!!! note "Approximation error"
- For the mass equation, the only approximation we have performed is
- quadrature. No interpolation or finite difference error is present.
+!!! note "Storage convention"
+ We use the column-major convention (Julia, MATLAB, Fortran), and not the
+ row-major convention (Python, C). Thus the ``x^1``-index ``i`` will vary for
+ one whole cycle in the vectors before the
+ ``x^2``-index ``j``, ``x^3`` index ``k``, and component-index ``\alpha``
+ are incremented, e.g. ``u_h = (u^1_{(1, 1, 1)},
+ u^1_{(2, 1, 1)}, \dots u^3_{(N_{u^3}(1), N_{u^3}(2), N_{u^3}(3))})`` in 3D.
-### Momentum equations
+## Fourth order accurate discretization
-Grouping the convection, pressure gradient, diffusion, and body force terms in
-each of their own integrals, we get, for all ``\mathcal{O} \subset \Omega``:
+The above discretization is second order accurate.
+A fourth order accurate discretization can be obtained by judiciously combining
+the second order discretization with itself on a grid with three times larger
+cells in each dimension [Verstappen2003](@cite) [Sanderse2014](@cite). The
+coarse discretization is identical, but the mass equation is derived for the
+three times coarser control volume
```math
-\frac{\partial }{\partial t} \int_\mathcal{O} u^\alpha \, \mathrm{d} \Omega
-=
-- \sum_{\beta = 1}^d \int_{\partial \mathcal{O}} u^\alpha u^\beta n^\beta \,
-\mathrm{d} \Gamma
-- \int_{\partial \mathcal{O}} p n^\alpha \, \mathrm{d} \Gamma
-+ \nu \sum_{\beta = 1}^d \int_{\partial \mathcal{O}} \frac{\partial
-u^\alpha}{\partial x^\beta} n^\beta \, \mathrm{d} \Gamma
-+ \int_\mathcal{O} f^\alpha \mathrm{d} \Omega,
+\Omega^3_I =
+\bigcup_{\alpha = 1}^d \Omega_{I - e_\alpha} \cup
+\Omega_I \cup \Omega_{I + e_\alpha},
```
-
-where ``n = (n^1, \dots, n^d)`` is the surface normal vector to ``\partial
-\mathcal{O}``.
-
-This time, we will not let ``\mathcal{O}`` be the reference finite volume
-``\Omega_{I}`` (the ``p``-volume), but rather the shifted ``u^\alpha``-volume.
-Setting ``\mathcal{O} = \Omega_{I + \delta(\alpha) / 2}`` (with right/rear/top
-``\beta``-faces ``\Gamma^\beta_{I + \delta(\alpha) / 2 + \delta(\beta) / 2}``)
-gives
+while the momentum equation is
+derived for its shifted variant ``\Omega^3_{I + h_\alpha}``.
+The resulting fourth order accurate equations are given by
```math
-\begin{split}
- \frac{\partial }{\partial t}
- \int_{\Omega_{I + \delta(\alpha) / 2}}
- \! \! \!
- \! \! \!
- \! \! \!
- \! \! \!
- u^\alpha \, \mathrm{d} \Omega
- =
- & -
- \sum_{\beta = 1}^d \left(
- \int_{\Gamma^{\beta}_{I + \delta(\alpha) / 2 + \delta(\beta) / 2}}
- \! \! \!
- \! \! \!
- \! \! \!
- \! \! \!
- u^\alpha u^\beta \, \mathrm{d} \Gamma
- - \int_{\Gamma^{\beta}_{I + \delta(\alpha) / 2 - \delta(\beta) / 2}}
- \! \! \!
- \! \! \!
- \! \! \!
- \! \! \!
- \! \! \!
- \! \! \!
- u^\alpha u^\beta \, \mathrm{d} \Gamma
- \right) \\
- & -
- \left(
- \int_{\Gamma^{\alpha}_{I + \delta(\alpha)}} p \, \mathrm{d} \Gamma
- - \int_{\Gamma^{\alpha}_{I}} p \, \mathrm{d} \Gamma
- \right) \\
- & + \nu \sum_{\beta = 1}^d
- \left(
- \int_{\Gamma^{\beta}_{I + \delta(\alpha) / 2 + \delta(\beta) / 2}}
- \frac{\partial u^\alpha}{\partial x^\beta} \, \mathrm{d} \Gamma
- - \int_{\Gamma^{\beta}_{I + \delta(\alpha) / 2 - \delta(\beta) / 2}}
- \frac{\partial u^\alpha}{\partial x^\beta} \, \mathrm{d} \Gamma
- \right) \\
- & +
- \int_{\Omega_{I + \delta(\alpha) / 2}}
- f^\alpha \, \mathrm{d} \Omega
-\end{split}
+\sum_{\alpha = 1}^d
+(\delta_\alpha u^\alpha)_I
+-
+\frac{| \Omega^3_I |}{3^{2 + d} | \Omega_I |}
+\sum_{\alpha = 1}^d
+(\delta^3_\alpha u^\alpha)_I
+= 0
```
-This equation is still exact. We now introduce some approximations on
-``\Omega_{I + \delta(\alpha) / 2}`` and its boundaries to remove all unknown
-continuous quantities.
-
-1. We replace the integrals with a mid-point quadrature rule.
-1. The mid-point values of derivatives are approximated using a central-like
- finite difference:
- ```math
- \frac{\partial u^\alpha}{\partial x^\beta}(x_I) \approx
- \frac{u^\alpha_{I + \delta(\beta) / 2}
- - u^\alpha_{I - \delta(\beta) / 2}}{x^\beta_{I(\beta) + 1 / 2} -
- x^\beta_{I(\beta) - 1 / 2}}.
- ```
-1. Quantities outside their canonical positions are obtained through
- interpolation.
-
-Finally, the discrete ``\alpha``-momentum equations are given by
+and
```math
-\begin{split}
- \left| \Omega_{I + \delta(\alpha) / 2} \right|
- & \frac{\mathrm{d} }{\mathrm{d} t} u^\alpha_{I + \delta(\alpha) / 2} = \\
- - & \sum_{\beta = 1}^d \left| \Gamma^\beta_{I + \delta(\alpha) / 2} \right|
- \left(
- (u^\alpha
- u^\beta)_{I + \delta(\alpha) / 2 + \delta(\beta) / 2}
- -
- (u^\alpha
- u^\beta )_{I + \delta(\alpha) / 2 - \delta(\beta) / 2}
- \right) \\
- - & \left| \Gamma^\alpha_{I + \delta(\alpha) / 2} \right|
- \left( p_{I + \delta(\alpha)} - p_{I} \right) \\
- + & \nu \sum_{\beta = 1}^d \left| \Gamma^\beta_{I + \delta(\alpha) / 2}
- \right|
- \left(
- \frac{u^\alpha_{I + \delta(\alpha) / 2 + \delta(\beta)} - u^\alpha_{I +
- \delta(\alpha) / 2}}{x^\beta_{I(\beta) + 1} - x^\beta_{I(\beta)}}
- - \frac{u^\alpha_{I + \delta(\alpha) / 2} - u^\alpha_{I +
- \delta(\alpha) / 2 - \delta(\beta)}}{x^\beta_{I(\beta)}
- - x^\beta_{I(\beta) - 1}}
- \right) \\
- + & \left| \Omega_{I + \delta(\alpha) / 2} \right| f^\alpha(x_{I +
- \delta(\alpha) / 2}).
-\end{split}
+\frac{\mathrm{d}}{\mathrm{d} t} u^\alpha_{I} =
+- \sum_{\beta = 1}^d
+(\delta_\beta (u^\alpha u^\beta))_{I}
++ \nu \sum_{\beta = 1}^d
+(\delta_\beta \delta_\beta u^\alpha)_{I}
++ f^\alpha(x_{I})
+- (\delta_\alpha p)_{I}
++ \text{fourth order},
```
+where
-## Boundary conditions
-
-Depending on the type of boundary conditions, certain indices used in the left-
-and right hand sides of the equations may not be part of the solution vectors,
-even if they are indeed at their canonical positions. Consider for example the
-``\alpha``-left/front/bottom boundary ``\{ x \in \Omega | x^{\alpha} =
-x^\alpha_{0} \}``. Let ``\Omega_I`` be one of the reference finite volumes
-touching this boundary, i.e. ``I(\alpha) = \frac{1}{2}``.
-
-- For periodic boundary conditions, we also consider the opposite boundary.
- We add two "ghost" reference volumes, one at each side of ``\Omega``:
- 1. The ghost volume ``\Omega_{I - \delta(\alpha)}`` has the same shape as
- ``\Omega_{I + (N(\alpha) - 1) \delta(\alpha)}`` and contains the same
- components:
- - ``u^\alpha_{I - \delta(\alpha) / 2} = u^\alpha_{I + (N(\alpha) - 1 / 2)
- \delta(\alpha)}``,
- - ``u^\beta_{I + \delta(\beta) / 2 - \delta(\alpha)} =
- u^\beta_{I + \delta(\beta) / 2 + N(\alpha) \delta(\alpha)}`` for ``\beta
- \neq \alpha``,
- - ``p_{I - \delta(\alpha)} = p_{I + (N(\alpha) - 1) \delta(\alpha)}``.
- 2. The ghost volume ``\Omega_{I + N(\alpha) \delta(\alpha)}`` has the same
- shape as ``\Omega_I`` and contains the same pressure and velocity
- components.
-- For Dirichlet boundary conditions, all the veloctity components are
- prescribed. For the normal velocity component, this is straightforward:
- ```math
- u^\alpha_{I - \delta(\alpha) / 2} = u^\alpha(x_{I - \delta(\alpha) / 2}).
- ```
- The parallel (``\beta \neq \alpha``) velocity components
- ``u^\beta_{I + \delta(\beta) / 2 - \delta(\alpha)}``
- appear in some of the right hand side expressions. Their ``\alpha``
- position ``x^\alpha_{- 1 / 2}`` has actually never been defined,
- so we simply define it to be on the boundary itself:
- ``x^\alpha_{- 1 / 2} = x^\alpha_0``. The value can then be
- prescribed:
- ```math
- u^\beta_{I + \delta(\beta) / 2 - \delta(\alpha)} = u^\beta \left( x_{I +
- \delta(\beta) / 2 - \delta(\alpha)} \right) = u^\beta \left( x_{I +
- \delta(\beta) / 2 - \delta(\alpha) / 2} \right).
- ```
- The pressure does not require any boundary modifications.
-- For a symmetric boundary, the normal component has zero Dirichlet
- conditions, while the parallel components have zero Neumann conditions. For
- this, we add a ghost volume ``\Omega_{I - \delta(\alpha)}`` which has the
- same shape as ``\Omega_{I}``, meaning that ``x^\alpha_{-1} = x^\alpha_0 -
- (x^\alpha_1 - x^\alpha_0) = - x^\alpha_1``. The pressure in this volume is
- never used, but we set ``u^\alpha_{I - \delta(\alpha) / 2} = 0`` and
- ``u^\beta_{I + \delta(\beta) / 2 - \delta(\alpha)} = u^\beta_{I +
- \delta(\beta) / 2}`` for ``\beta \neq \alpha``.
-- For a pressure boundary, the value of the pressure is prescribed, while the
- velocity has zero Neumann boundary conditions. For this, we add an
- infinitely thin ghost volume ``\Omega_{I - \delta(\alpha)}``, by setting
- ``x^\alpha_{-1} = x^\alpha_0 - \epsilon`` for some epsilon. We then let this
- thickness go to zero *after* substituting the boundary conditions in
- the momentum equations. The pressure is prescribed by
- ```math
- p_{I - \delta(\alpha)} = p(x_{I - \delta(\alpha)}) \underset{\epsilon \to 0}{\to} p(x_{I - \delta(\alpha) / 2}).
- ```
- The Neumann boundary conditions are obtained by setting
- ``u^\alpha_{I - 3 / 2 \delta(\alpha)} = u^\alpha_{I - \delta(\alpha) / 2}``
- and
- ``u^\beta_{I - \delta(\alpha) + \delta(\beta) / 2} = u^\beta_{I +
- \delta(\beta) / 2}`` for ``\beta \neq \alpha``.
- Note that the normal velocity component at the boundary ``u^\alpha_{I -
- \delta(\alpha) / 2}`` is now a degree of freedom, and we need to
- observe that the first normal derivative in the diffusion term of its
- corresponding momentum equation is zero before taking the limit as ``\epsilon
- \to 0``, to avoid dividing by zero.
-
-It should now be clear from the above cases which components of the discrete
-velocity and pressure are unknown degrees of freedom, and which components are
-prescribed or obtained otherwise. The *unknown* degrees of freedom are stored
-in the vectors ``u_h = (u^1_h, \dots, u^d_h)`` and ``p_h`` using the
-column-major convention. Note that the ``d`` discrete velocity fields ``u^1_h,
-\dots u^d_h`` may have different numbers of elements.
-
-!!! note "Storage convention"
- We use the column-major convention (Julia, MATLAB, Fortran), and not the
- row-major convention (Python, C). Thus the ``x^1``-index ``i`` will vary for
- one whole cycle in the vectors before the
- ``x^2``-index ``j``, ``x^3`` index ``k``, and component-index ``\alpha``
- are incremented, e.g. ``u_h = (u^1_{(1, 1, 1)},
- u^1_{(2, 1, 1)}, \dots u^3_{(N_{u^3}(1), N_{u^3}(2), N_{u^3}(3))})`` in 3D.
-
+```math
+(\delta^3_\alpha \varphi)_I =
+\frac{\varphi_{I + 3 h_\alpha} -
+\varphi_{I - 3 h_\alpha}}{\Delta^\alpha_{I_\alpha - 1} +
+\Delta^\alpha_{I_\alpha} + \Delta^\alpha_{I_\alpha + 1}}.
+```
## Matrix representation
-We can write the mass and momentum equations in matrix form. The discrete mass
-equation then becomes
+We can write the mass and momentum equations in matrix form. We will use the
+same matrix notation for the second- and fourth order accurate discretizations.
+The discrete mass equation becomes
+
```math
M u_h + y_M = 0,
```
+
where ``M`` is the discrete divergence operator and ``y_M`` contains the
boundary value contributions of the velocity to the divergence field.
@@ -399,19 +181,23 @@ The discrete momentum equations become
\end{split}
```
-where ``C`` is the convection operator (including boundary contributions),
+where ``C`` is she convection operator (including boundary contributions),
``D`` is the diffusion operator, ``y_D`` is boundary contribution to the
-diffusion term, ``G = \Omega_h^{-1} M^\mathsf{T}`` is the pressure gradient
-operator, ``y_G`` contains the boundary contribution of the pressure to the
-pressure gradient (only non-zero for pressure boundary conditions), and
-``\Omega_h`` is a diagonal matrix containing the velocity volumes. The term
-``F`` refers to all the forces except for the pressure gradient.
+diffusion term,
+``G = W_u^{-1} M^\mathsf{T} W`` is the pressure gradient
+operator,
+``y_G`` contains the boundary contribution of the pressure to the
+pressure gradient (only non-zero for pressure boundary conditions),
+``W_u`` is a diagonal matrix containing the velocity volume sizes
+``| \Omega_{I + \delta(\alpha) / 2} |``, and ``W`` is a diagonal matrix
+containing the reference volume sizes ``| \Omega_I |``.
+The term ``F`` refers to all the forces except for the pressure gradient.
!!! note "Volume normalization"
- All the operators (except for ``M``) have been divided by the velocity
- volume sizes ``\Omega_h``. As a result, the operators have the same units
- as their continuous counterparts.
+ All the operators have been divided by the velocity volume sizes.
+ As a result, the operators have the same units as their
+ continuous counterparts.
## Discrete pressure Poisson equation
@@ -423,11 +209,11 @@ discrete divergence operator ``M`` to the discrete momentum equations yields
the discrete pressure Poisson equation
```math
-L p_h = M (F(u_h) - y_G) + \frac{\mathrm{d} y_M}{\mathrm{d} t},
+L p_h = W M (F(u_h) - y_G) + W \frac{\mathrm{d} y_M}{\mathrm{d} t},
```
-where ``L = M G`` is a discrete Laplace operator. It is positive
-symmetric since ``G = \Omega_h^{-1} M^\mathsf{T}``.
+where ``L = W M G = W M W_u^{-1} M^\mathsf{T} W`` is a discrete Laplace
+operator. It is positive symmetric.
!!! note "Unsteady Dirichlet boundary conditions"
@@ -455,17 +241,17 @@ symmetric since ``G = \Omega_h^{-1} M^\mathsf{T}``.
discrete Poisson equation:
```math
- p_h = L^{-1} M (F(u_h) - y_G) + L^{-1} \frac{\mathrm{d} y_M}{\mathrm{d} t}.
+ p_h = L^{-1} W M (F(u_h) - y_G) + L^{-1} W \frac{\mathrm{d} y_M}{\mathrm{d} t}.
```
The momentum equations then become
```math
- \frac{\mathrm{d} u_h}{\mathrm{d} t} = (I - G L^{-1} M)
- (F(u_h) - y_G) - G L^{-1} \frac{\mathrm{d} y_M}{\mathrm{d} t}.
+ \frac{\mathrm{d} u_h}{\mathrm{d} t} = (I - G L^{-1} W M)
+ (F(u_h) - y_G) - G L^{-1} W \frac{\mathrm{d} y_M}{\mathrm{d} t}.
```
- The matrix ``(I - G L^{-1} M)`` is a projector onto the space
+ The matrix ``(I - G L^{-1} W M)`` is a projector onto the space
of discretely divergence free velocities. However, using this formulation
would require an efficient way to perform the projection without assembling
the operator matrix ``L^{-1}``, which would be very costly.
@@ -482,69 +268,20 @@ the squares.
### Vorticity
-In 2D, the vorticity is a scalar. Integrating the vorticity ``\omega =
--\frac{\partial u^1}{\partial x^2} + \frac{\partial u^2}{\partial x^1}`` over
-the vorticity volume ``\Omega_{I + \delta(1) / 2 + \delta(2) / 2}`` gives
-
-```math
-\begin{split}
-\int_{\Omega_{I + \delta(1) / 2 + \delta(2) / 2}} \omega \, \mathrm{d} \Omega
-= & - \left(
-\int_{\Gamma^2_{I + \delta(1) / 2 + \delta(2)}} u^1 \, \mathrm{d} \Gamma
-- \int_{\Gamma^2_{I + \delta(1) / 2}} u^1 \, \mathrm{d} \Gamma
-\right) \\
-& + \left(
-\int_{\Gamma^1_{I + \delta(1) + \delta(2) / 2}} u^2 \, \mathrm{d} \Gamma
-- \int_{\Gamma^1_{I + \delta(2) / 2}} u^2 \, \mathrm{d} \Gamma
-\right)
-\end{split}.
-```
-
-Using quadrature, the discrete vorticity in the corner is given by
+In 2D, the vorticity is a scalar. We define it as
```math
-\begin{split}
-\left| \Omega_{I + \delta(1) / 2 + \delta(2) / 2} \right|
-\omega_{I + \delta(1) / 2 + \delta(2) / 2} =
-& - \left| \Gamma^2_{I + \delta(1) / 2 + \delta(2) / 2} \right|
-(u^1_{I + \delta(1) / 2 + \delta(2)}
-- u^1_{I + \delta(1) / 2}) \\
-& + \left| \Gamma^1_{I + \delta(1) / 2 + \delta(2) / 2} \right|
-(u^2_{I + \delta(1) + \delta(2) / 2}
-- u^2_{I + \delta(2) / 2})
-\end{split}.
+\omega = -\delta^2 u^1 + \delta^1 u^2.
```
The 3D vorticity is a vector field ``(\omega^1, \omega^2, \omega^3)``.
Noting ``\alpha^+ = \operatorname{mod}_3(\alpha + 1)`` and
-``\alpha^- = \operatorname{mod}_3(\alpha - 1)``, the vorticity is obtained
+``\alpha^- = \operatorname{mod}_3(\alpha - 1)``, the vorticity is defined as
through
```math
-\begin{split}
-\int_{\Omega_{I + \delta(\alpha^+) / 2 + \delta(\alpha^-) / 2}} \omega \, \mathrm{d} \Omega
-= & - \left(
-\int_{\Gamma^{\alpha^-}_{I + \delta(\alpha^+) / 2 + \delta(\alpha^-)}} u^{\alpha^+} \, \mathrm{d} \Gamma
-- \int_{\Gamma^{\alpha^-}_{I + \delta(\alpha^+) / 2}} u^{\alpha^+} \, \mathrm{d} \Gamma
-\right) \\
-& + \left(
-\int_{\Gamma^{\alpha^+}_{I + \delta(\alpha^+) + \delta(\alpha^-) / 2}} u^{\alpha^-} \, \mathrm{d} \Gamma
-- \int_{\Gamma^{\alpha^+}_{I + \delta(\alpha^-) / 2}} u^{\alpha^-} \, \mathrm{d} \Gamma
-\right)
-\end{split}.
-```
-
-```math
-\begin{split}
-\left| \Omega_{I + \delta(\alpha^+) / 2 + \delta(\alpha^-) / 2} \right|
-\omega_{I + \delta(\alpha^+) / 2 + \delta(\alpha^-) / 2} =
-& - \left| \Gamma^{\alpha^-}_{I + \delta(\alpha^+) / 2 + \delta(\alpha^-) / 2} \right|
-(u^{\alpha^+}_{I + \delta(\alpha^+) / 2 + \delta(\alpha^-)}
-- u^{\alpha^+}_{I + \delta(\alpha^+) / 2}) \\
-& + \left| \Gamma^{\alpha^+}_{I + \delta(\alpha^+) / 2 + \delta(\alpha^-) / 2} \right|
-(u^{\alpha^-}_{I + \delta(\alpha^+) + \delta(\alpha^-) / 2}
-- u^{\alpha^-}_{I + \delta(\alpha^-) / 2})
-\end{split}.
+\omega^\alpha = - \delta^{\alpha^-} u^{\alpha^+} +
+\delta^{\alpha^+} u^{\alpha^-}.
```
## Stream function
@@ -555,16 +292,16 @@ yields
```math
\begin{split}
-- \int_{\Omega_{I + \delta(1) / 2 + \delta(2) / 2}} \omega \, \mathrm{d} \Omega
-& = \int_{\Omega_{I + \delta(1) / 2 + \delta(2) / 2}} \nabla^2 \psi \,
+- \int_{\Omega_{I + h_1 + h_2}} \omega \, \mathrm{d} \Omega
+& = \int_{\Omega_{I + h_1 + h_2}} \nabla^2 \psi \,
\mathrm{d} \Omega \\
-& = \int_{\Gamma^1_{I + \delta(1) + \delta(2) / 2}} \frac{\partial \psi}{\partial x^1}
+& = \int_{\Gamma^1_{I + e_1 + h_2}} \frac{\partial \psi}{\partial x^1}
\, \mathrm{d} \Gamma
-- \int_{\Gamma^1_{I + \delta(2) / 2}} \frac{\partial \psi}{\partial x^1}
+- \int_{\Gamma^1_{I + h_2}} \frac{\partial \psi}{\partial x^1}
\, \mathrm{d} \Gamma \\
-& + \int_{\Gamma^2_{I + \delta(1) / 2 + \delta(2)}} \frac{\partial \psi}{\partial x^2}
+& + \int_{\Gamma^2_{I + h_1 + e_2}} \frac{\partial \psi}{\partial x^2}
\, \mathrm{d} \Gamma
-- \int_{\Gamma^2_{I + \delta(1) / 2}} \frac{\partial \psi}{\partial x^2}
+- \int_{\Gamma^2_{I + h_1}} \frac{\partial \psi}{\partial x^2}
\, \mathrm{d} \Gamma.
\end{split}
```
@@ -575,17 +312,17 @@ equation for the stream function at the vorticity point:
```math
\begin{split}
-\left| \Gamma^1_{I + \delta(1) / 2 + \delta(2) / 2} \right|
+\left| \Gamma^1_{I + h_1 + h_2} \right|
\left(
- \frac{\psi_{I + 3 / 2 \delta(1) + \delta(2) / 2} - \psi_{I + \delta(1) / 2 + \delta(2) / 2}}{x^1_{I(1) + 3 / 2} - x^1_{I(1) + 1 /2}}
-- \frac{\psi_{I + \delta(1) / 2 + \delta(2) / 2} - \psi_{I - \delta(1) / 2 + \delta(2) / 2}}{x^1_{I(1) + 1 / 2} - x^1_{I(1) - 1 / 2}}
+ \frac{\psi_{I + 3 / h_1 + h_2} - \psi_{I + h_1 + h_2}}{x^1_{I_1 + 3 / 2} - x^1_{I_1 + 1 /2}}
+- \frac{\psi_{I + h_1 + h_2} - \psi_{I - h_1 + h_2}}{x^1_{I_1 + 1 / 2} - x^1_{I_1 - 1 / 2}}
\right) & + \\
-\left| \Gamma^2_{I + \delta(1) / 2 + \delta(2) / 2} \right|
+\left| \Gamma^2_{I + h_1 + h_2} \right|
\left(
-\frac{\psi_{I + \delta(1) / 2 + 3 / 2 \delta(2)} - \psi_{I + \delta(1) / 2 + \delta(2) / 2}}{x^2_{I(1) + 3 / 2} - x^2_{I(1) + 1 / 2}}
--\frac{\psi_{I + \delta(1) / 2 + \delta(2) / 2} - \psi_{I + \delta(1) / 2 - \delta(2) / 2}}{x^2_{I(2) + 1 / 2} - x^2_{I(2) - 1 / 2}}
+\frac{\psi_{I + h_1 + 3 h_2} - \psi_{I + h_1 + h_2}}{x^2_{I_1 + 3 / 2} - x^2_{I_1 + 1 / 2}}
+-\frac{\psi_{I + h_1 + h_2} - \psi_{I + h_1 - h_2}}{x^2_{I_2 + 1 / 2} - x^2_{I_2 - 1 / 2}}
\right) & = \\
-\left| \Omega_{I + \delta(1) / 2 + \delta(2) / 2} \right|
-\omega_{I + \delta(1) / 2 + \delta(2) / 2} &
+\left| \Omega_{I + h_1 + h_2} \right|
+\omega_{I + h_1 + h_2} &
\end{split}
```
diff --git a/docs/src/equations/time.md b/docs/src/equations/time.md
index ae39eac7c..a8eef08d8 100644
--- a/docs/src/equations/time.md
+++ b/docs/src/equations/time.md
@@ -1,16 +1,20 @@
+```@meta
+CurrentModule = IncompressibleNavierStokes
+```
+
# Time discretization
The spatially discretized Navier-Stokes equations form a differential-algebraic
system, with an ODE for the velocity
```math
-\frac{\mathrm{d} u_h}{\mathrm{d} t} = F(u_h, t) - (G p_h + y_G)
+\frac{\mathrm{d} u}{\mathrm{d} t} = F(u, t) - (G p + y_G)
```
subject to the algebraic constraint formed by the mass equation
```math
-M u_h + y_M = 0.
+M u + y_M = 0.
```
In the end of the [previous section](spatial.md), we differentiated the mass
@@ -18,7 +22,7 @@ equation in time to obtain a discrete pressure Poisson equation. This equation
includes the term ``\frac{\mathrm{d} y_M}{\mathrm{d} t}``, which is non-zero if
an unsteady flow of mass is added to the domain (Dirichlet boundary
conditions). This term ensures that the time-continuous discrete velocity field
-``u_h(t)`` stays divergence free (conserves mass). However, if we directly
+``u(t)`` stays divergence free (conserves mass). However, if we directly
discretize this system in time, the mass preservation may actually not be
respected. For this, we will change the definition of the pressure such that
the time-discretized velocity field is divergence free at each time step and
@@ -27,79 +31,106 @@ each time sub-step (to be defined in the following).
Consider the interval ``[0, T]`` for some simulation time ``T``. We will divide
it into ``N`` sub-intervals ``[t^n, t^{n + 1}]`` for ``n = 0, \dots, N - 1``,
with ``t^0 = 0``, ``t^N = T``, and increment ``\Delta t^n = t^{n + 1} - t^n``.
-We define ``U^n \approx u_h(t^n)`` as an approximation to the exact discrete
-velocity field ``u_h(t^n)``, with ``U^0 = u_h(0)`` starting from the exact
+We define ``u^n \approx u(t^n)`` as an approximation to the exact discrete
+velocity field ``u(t^n)``, with ``u^0 = u(0)`` starting from the exact
initial conditions. We say that the time integration scheme (definition of
-``U^n``) is accurate to the order ``r`` if ``U^n = u_h(t^n) +
+``u^n``) is accurate to the order ``r`` if ``u^n = u(t^n) +
\mathcal{O}(\Delta t^r)`` for all ``n``.
+IncompressibleNavierStokes provides a collection of explicit and implicit
+[Runge-Kutta methods](../api/tableaux.md), in addition to Adams-Bashforth
+Crank-Nicolson and one-leg beta method time steppers.
+
+The code is currently not adapted to time steppers from
+[DifferentialEquations.jl](https://docs.sciml.ai/DiffEqDocs/stable/solvers/dae_solve/),
+but they may be integrated in the future.
+
+```@docs
+AbstractODEMethod
+AbstractRungeKuttaMethod
+isexplicit
+lambda_conv_max
+lambda_diff_max
+ode_method_cache
+nstage
+runge_kutta_method
+timestep
+timestep!
+```
+
## Explicit Runge-Kutta methods
See Sanderse [Sanderse2012](@cite).
-Consider the velocity field ``U_0`` at a certain time ``t_0``. We will now
+Consider the velocity field ``u_0`` at a certain time ``t_0``. We will now
perform one time step to ``t = t_0 + \Delta t``. For explicit Runge-Kutta
methods, this time step is divided into ``s`` sub-steps ``t_i = t_0 + \Delta
t_i`` with increment ``\Delta t_i = c_i \Delta t``. The final substep performs
the full time step ``\Delta t_s = \Delta t`` such that ``t_s = t``.
-For ``i = 1, \dots, s``, the intermediate velocity ``U_i`` and pressure ``P_i``
+For ``i = 1, \dots, s``, the intermediate velocity ``u_i`` and pressure ``p_i``
are computed as follows:
```math
\begin{split}
-F_i & = F(U_{i - 1}, t_{i - 1}) - y_G(t_{i - 1}) \\
-V_i & = U_0 + \Delta t \sum_{j = 1}^i a_{i j} F_j \\
-L P_i & = \frac{1}{c_i} \sum_{j = 1}^i a_{i j} F_j +
-\frac{y_M(t_i) - y_M(t_0)}{\Delta t_i} \\
-& = \frac{(M V_i + y_M(t_i)) - (M U_0 + y_M(t_0))}{\Delta t_i^n} \\
-& = \frac{M V_i + y_M(t_i)}{\Delta t_i^n} \\
-U_i & = V_i - \Delta t_i G P_i,
+k_i & = F(u_{i - 1}, t_{i - 1}) - y_G(t_{i - 1}) \\
+v_i & = u_0 + \Delta t \sum_{j = 1}^i a_{i j} k_j \\
+L p_i & = W M \frac{1}{c_i} \sum_{j = 1}^i a_{i j} k_j +
+W \frac{y_M(t_i) - y_M(t_0)}{\Delta t_i} \\
+& = W \frac{(M v_i + y_M(t_i)) - (M u_0 + y_M(t_0))}{\Delta t_i^n} \\
+& = W \frac{M v_i + y_M(t_i)}{\Delta t_i^n} \\
+u_i & = v_i - \Delta t_i G p_i,
\end{split}
```
where ``(a_{i j})_{i j}`` are the Butcher tableau coefficients of the
RK-method, with the convention ``c_i = \sum_{j = 1}^i a_{i j}``.
-Finally, we set ``U = U_s``. If ``U_0 = u_h(t_0)``, we get the accuracy ``U =
-u_h(t) + \mathcal{O}(\Delta t^{r + 1})``, where ``r`` is the order of the
+Finally, we return ``u_s``. If ``u_0 = u(t_0)``, we get the accuracy ``u_s =
+u(t) + \mathcal{O}(\Delta t^{r + 1})``, where ``r`` is the order of the
RK-method. If we perform ``n`` RK time steps instead of one, starting at exact
-initial conditions ``U^0 = u_h(0)``, then ``U^n = u_h(t^n) + \mathcal{O}(\Delta
-t^r)`` for all ``n \in \{1, \dots, N\}``. Note that for a given ``U``, the
-corresponding pressure ``P`` can be calculated to the same accuracy as ``U`` by
+initial conditions ``u^0 = u(0)``, then ``u^n = u(t^n) + \mathcal{O}(\Delta
+t^r)`` for all ``n \in \{1, \dots, N\}``. Note that for a given ``u``, the
+corresponding pressure ``p`` can be calculated to the same accuracy as ``u`` by
doing an additional pressure projection after each outer time step ``\Delta t``
(if we know ``\frac{\mathrm{d} y_M}{\mathrm{d} t}(t)``), or to first order
-accuracy by simply taking ``P = P_s``.
+accuracy by simply returning ``p_s``.
-Note that each of the sub-step velocities ``U_i`` is divergence free, after
-projecting the tentative velocities ``V_i``. This is ensured due to the
+Note that each of the sub-step velocities ``u_i`` is divergence free, after
+projecting the tentative velocities ``v_i``. This is ensured due to the
judiciously chosen replacement of ``\frac{\mathrm{d} y_M}{\mathrm{d} t}(t_i)``
with ``(y_M(t_i) - y_M(t_0)) / \Delta t_i``. The space-discrete
divergence-freeness is thus perfectly preserved, even though the time
discretization introduces other errors.
+```@docs
+ExplicitRungeKuttaMethod
+```
## Implicit Runge-Kutta methods
See Sanderse [Sanderse2013](@cite).
+```@docs
+ImplicitRungeKuttaMethod
+```
## Adams-Bashforth Crank-Nicolson method
We here require that the time step ``\Delta t`` is constant. This methods uses
Adams-Bashforth for the convective terms and Crank-Nicolson stepping for the
-diffusion and body force terms. Given the velocity field ``U_0 = u_h(t_0)`` at
-a time ``t_0`` and its previous value ``U_{-1} = u_h(t_0 - \Delta t)`` at the
-previous time ``t_{-1} = t_0 - \Delta t``, the predicted velocity field ``U``
+diffusion and body force terms. Given the velocity field ``u_0 = u(t_0)`` at
+a time ``t_0`` and its previous value ``u_{-1} = u(t_0 - \Delta t)`` at the
+previous time ``t_{-1} = t_0 - \Delta t``, the predicted velocity field ``u``
at the time ``t = t_0 + \Delta t`` is defined by first computing a tentative
velocity:
```math
\begin{split}
-\frac{V - U_0}{\Delta t}
-& = - (\alpha_0 C(U_0, t_0) + \alpha_{-1} C(U_{-1}, t_{-1})) \\
-& + \theta (D U_0 + y_D(t_0)) + (1 - \theta) (D V + y_D(t)) \\
+\frac{v - u_0}{\Delta t}
+& = - (\alpha_0 C(u_0, t_0) + \alpha_{-1} C(u_{-1}, t_{-1})) \\
+& + \theta (D u_0 + y_D(t_0)) + (1 - \theta) (D v + y_D(t)) \\
& + \theta f(t_0) + (1 - \theta) f(t) \\
& - (G p_0 + y_G(t_0)),
\end{split}
@@ -108,74 +139,81 @@ velocity:
where ``\theta \in [0, 1]`` is the Crank-Nicolson parameter (``\theta =
\frac{1}{2}`` for second order convergence), ``(\alpha_0, \alpha_{-1}) = \left(
\frac{3}{2}, -\frac{1}{2} \right)`` are the Adams-Bashforth coefficients, and
-``V`` is a tentative velocity yet to be made divergence free. We can group the
-terms containing ``V`` on the left hand side, to obtain
+``v`` is a tentative velocity yet to be made divergence free. We can group the
+terms containing ``v`` on the left hand side, to obtain
```math
\begin{split}
-\left( \frac{1}{\Delta t} I - (1 - \theta) D \right) V
-& = \left(\frac{1}{\Delta t} I - \theta D \right) U_0 \\
-& - (\alpha_0 C(U_0, t_0) + \alpha_{-1} C(U_{-1}, t_{-1})) \\
+\left( \frac{1}{\Delta t} I - (1 - \theta) D \right) v
+& = \left(\frac{1}{\Delta t} I - \theta D \right) u_0 \\
+& - (\alpha_0 C(u_0, t_0) + \alpha_{-1} C(u_{-1}, t_{-1})) \\
& + \theta y_D(t_0) + (1 - \theta) y_D(t) \\
& + \theta f(t_0) + (1 - \theta) f(t) \\
-& - (G P_0 + y_G(t_0)).
+& - (G p_0 + y_G(t_0)).
\end{split}
```
-We can compute ``V`` by inverting the positive definite matrix ``\left(
+We can compute ``v`` by inverting the positive definite matrix ``\left(
\frac{1}{\Delta t} I - \theta D \right)`` for the given right hand side using a
suitable linear solver. Assuming ``\Delta t`` is constant, we can precompute a
Cholesky factorization of this matrix before starting time stepping.
-We then compute the pressure difference ``\Delta P`` by solving
+We then compute the pressure difference ``\Delta p`` by solving
```math
-L \Delta P = \frac{M V + y_M(t)}{\Delta t} - M (y_G(t) - y_G(t_0)),
+L \Delta p = W \frac{M v + y_M(t)}{\Delta t} - W M (y_G(t) - y_G(t_0)),
```
-after which a divergence free velocity ``U`` can be enforced:
+after which a divergence free velocity ``u`` can be enforced:
```math
-U = V - \Delta t (G \Delta P + y_G(t) - y_G(t_0)).
+u = v - \Delta t (G \Delta p + y_G(t) - y_G(t_0)).
```
-A first order accurate prediction of the corresponding pressure is ``P = P_0 +
-\Delta P``. However, since this pressure is reused in the next time step, we
+A first order accurate prediction of the corresponding pressure is ``p = p_0 +
+\Delta p``. However, since this pressure is reused in the next time step, we
perform an additional pressure solve to avoid accumulating first order errors.
-The resulting pressure ``P`` is then accurate to the same order as ``U``.
+The resulting pressure ``p`` is then accurate to the same order as ``u``.
+```@docs
+AdamsBashforthCrankNicolsonMethod
+```
## One-leg beta method
See Verstappen and Veldman [Verstappen2003](@cite) [Verstappen1997](@cite).
We here require that the time step ``\Delta t`` is constant. Given the velocity
-``U_0`` and pressure ``P_0`` at the current time ``t_0`` and their previous
-values ``U_{-1}`` and ``P_{-1}`` at the time ``t_{-1} = t_0 - \Delta t``, we
-start by computing the "offstep" values ``V = (1 + \beta) V_0 - \beta V_{-1}``
-and ``Q = (1 + \beta) P_0 - \beta P_{-1}`` for some ``\beta = \frac{1}{2}``.
+``u_0`` and pressure ``p_0`` at the current time ``t_0`` and their previous
+values ``u_{-1}`` and ``p_{-1}`` at the time ``t_{-1} = t_0 - \Delta t``, we
+start by computing the "offstep" values ``v = (1 + \beta) v_0 - \beta v_{-1}``
+and ``Q = (1 + \beta) p_0 - \beta p_{-1}`` for some ``\beta = \frac{1}{2}``.
-A tentative velocity field ``W`` is then computed as follows:
+A tentative velocity field ``\tilde{v}`` is then computed as follows:
```math
-W = \frac{1}{\beta + \frac{1}{2}} \left( 2 \beta U_0 - \left( \beta -
-\frac{1}{2} \right) U_{-1} + \Delta t F(V, t) - \Delta t
+\tilde{v} = \frac{1}{\beta + \frac{1}{2}} \left( 2 \beta u_0 - \left( \beta -
+\frac{1}{2} \right) u_{-1} + \Delta t F(v, t) - \Delta t
(G Q + y_G(t)) \right).
```
-A pressure correction ``\Delta P `` is obtained by solving the Poisson equation
+A pressure correction ``\Delta p `` is obtained by solving the Poisson equation
```math
-L \Delta P = \frac{\beta + \frac{1}{2}}{\Delta t} (M W + y_M(t)).
+L \Delta p = \frac{\beta + \frac{1}{2}}{\Delta t} W (M \tilde{v} + y_M(t)).
```
Finally, the divergence free velocity field is given by
```math
-U = W - \frac{\Delta t}{\beta + \frac{1}{2}} G \Delta P,
+u = \tilde{v} - \frac{\Delta t}{\beta + \frac{1}{2}} G \Delta p,
```
while the second order accurate pressure is given by
```math
-P = 2 P_0 - P_{-1} + \frac{4}{3} \Delta P.
+p = 2 p_0 - p_{-1} + \frac{4}{3} \Delta p.
+```
+
+```@docs
+OneLegMethod
```
diff --git a/docs/src/features/bc.md b/docs/src/features/bc.md
index 9ee4f12b5..fcb244076 100644
--- a/docs/src/features/bc.md
+++ b/docs/src/features/bc.md
@@ -1,8 +1,21 @@
+```@meta
+CurrentModule = IncompressibleNavierStokes
+```
+
# Boundary conditions
-Various boundary conditions are supported. They are indicated by the symbols
+Each boundary has exactly one type of boundary conditions. For periodic
+boundary conditions, the opposite boundary must also be periodic.
+The available boundary conditions are given below.
+
+```@docs
+PeriodicBC
+DirichletBC
+SymmetricBC
+PressureBC
+```
-- `:periodic`
-- `:dirichlet`
-- `:symmetric`
-- `:pressure`
+```@docs
+offset_p
+offset_u
+```
diff --git a/docs/src/features/closure.md b/docs/src/features/closure.md
new file mode 100644
index 000000000..f2ec0f84f
--- /dev/null
+++ b/docs/src/features/closure.md
@@ -0,0 +1,117 @@
+# Neural closure models
+
+!!! note "`NeuralClorusure`"
+ These features are experimental, and require cloning
+ IncompressibleNavierStokes from GitHub:
+ ```sh
+ git clone https://github.com/agdestein/IncompressibleNavierStokes.jl
+ cd IncompressibleNavierStokes/lib/NeuralClosure
+ ```
+
+For [large eddy simulation (LES)](../features/les.md), a closure model is
+required. With IncompressibleNavierStokes, a neural closure model can be
+trained on filtered DNS data. The discrete DNS equations are given by
+
+```math
+\begin{split}
+M u & = 0, \\
+\frac{\mathrm{d} u}{\mathrm{d} t} & = F(u) - G p.
+\end{split}
+```
+
+Applying a spatial filter ``\Phi``, the extracted large scale components ``\bar{u} = \Phi u`` are governed by the equation
+
+```math
+\begin{split}
+M \bar{u} & = 0, \\
+\frac{\mathrm{d} \bar{u}}{\mathrm{d} t} & = F(\bar{u}) + c - G \bar{p},
+\end{split}
+```
+
+where the discretizations ``M``, ``F``, and ``G`` are adapted to the size of
+their inputs and ``c = \overline{F(u)} - F(\bar{u})`` is a commutator error. We
+here assumed that ``M`` and ``\Phi`` commute, which is the case for face
+averaging filters. Replacing ``c`` with a parameterized closure model
+``m(\bar{u}, \theta) \approx c`` gives the LES equations for the approximate
+large scale velocity ``\bar{v} \approx \bar{u}``
+
+```math
+\begin{split}
+M \bar{v} & = 0, \\
+\frac{\mathrm{d} \bar{v}}{\mathrm{d} t} & = F(\bar{v}) + m(\bar{v}, \theta) - G \bar{q}.
+\end{split}
+```
+
+## NeuralClosure module
+
+IncompressibleNavierStokes provides a NeuralClosure module.
+
+```@docs
+NeuralClosure.NeuralClosure
+```
+
+The following filters are available:
+
+```@docs
+NeuralClosure.FaceAverage
+NeuralClosure.VolumeAverage
+NeuralClosure.reconstruct!
+```
+
+The following functions create data.
+
+```@docs
+NeuralClosure.filtersaver
+NeuralClosure.create_les_data
+NeuralClosure.create_io_arrays
+```
+
+## Training
+
+To improve the model parameters, we exploit exact filtered DNS data ``\bar{u}``
+and exact commutator errors ``c`` obtained through DNS. The model is trained by
+minimizing the a priori loss function
+
+```math
+L^\text{prior}(\theta) = \| m(\bar{u}, \theta) - c \|^2,
+```
+
+or the a posteriori loss function
+
+```math
+L^\text{post}(\theta) = \| \bar{v}_\theta - \bar{u} \|^2,
+```
+
+where ``\bar{v}_\theta`` is the solution to the LES equation for the given
+parameters ``\theta``. The prior loss is easy to evaluate and easy to
+differentiate, as it does not involve solving the ODE. However, minimizing
+``L^\text{prior}`` does not take into account the effect of the prediction
+error on the LES solution error. The posterior loss does, but has a longer
+computational chain involving solving the LES ODE.
+
+```@docs
+NeuralClosure.create_dataloader_prior
+NeuralClosure.create_dataloader_post
+NeuralClosure.train
+NeuralClosure.create_loss_prior
+NeuralClosure.create_relerr_prior
+NeuralClosure.mean_squared_error
+NeuralClosure.create_loss_post
+NeuralClosure.create_relerr_post
+NeuralClosure.create_callback
+```
+
+## Neural architectures
+
+We provide two neural architectures: A convolutional neural network (CNN) and a Fourier neural operator (FNO).
+
+```@docs
+NeuralClosure.wrappedclosure
+NeuralClosure.create_closure
+NeuralClosure.create_tensorclosure
+NeuralClosure.collocate
+NeuralClosure.decollocate
+NeuralClosure.cnn
+NeuralClosure.fno
+NeuralClosure.FourierLayer
+```
diff --git a/docs/src/features/gpu.md b/docs/src/features/gpu.md
index d6a586769..b3f6ac5a2 100644
--- a/docs/src/features/gpu.md
+++ b/docs/src/features/gpu.md
@@ -1,33 +1,28 @@
# GPU Support
-If an Nvidia GPU is available, the default CPU solve call
-
-```julia
-solve_unsteady(setup, V₀, p₀, tlims; kwargs...)
-```
-
-can now be replaced with the following:
+IncompressibleNavierStokes supports various array types. The desired array type
+only has to be passed to the [`Setup`](@ref) function:
```julia
using CUDA
-solve_unsteady(
- setup, V₀, p₀, tlims;
- device = cu,
- kwargs...
-)
+setup = Setup(x...; kwargs..., ArrayType = CuArray)
```
-This moves the arrays and sparse operators to the GPU, outsourcing all array operations to the GPU.
-
-Limitations:
+All operators have been
+made are backend agnostic by using
+[KernelAbstractions.jl](https://github.com/JuliaGPU/KernelAbstractions.jl/).
+Even if a GPU is not available, the operators are multithreaded if
+Julia is started with multiple threads (e.g. `julia -t 4`)
-- [`DirectPressureSolver`](@ref) is currently not supported on the GPU. Use [`CGPressureSolver`](@ref) instead.
-- Unsteady boundary conditions are currently not supported on the GPU.
-- The code uses sparse matrices for discretization. For finer grids, these can take up a lot of memory on the GPU.
+- This has been tested with CUDA compatible GPUs.
- This has not been tested with other GPU interfaces, such as
- [AMDGPU.jl](https://github.com/JuliaGPU/AMDGPU.jl)
- [Metal.jl](https://github.com/JuliaGPU/Metal.jl)
- [oneAPI.jl](https://github.com/JuliaGPU/oneAPI.jl)
If they start supporting sparse matrices and fast Fourier transforms they
- could also be used. Alternatively, IncompressibleNavierStokes may also be
- refactored to apply the operators without assembling any sparse arrays.
+ could also be used.
+
+!!! note "`psolver_direct` on CUDA"
+ To use a specialized linear solver for CUDA, make sure to install and
+ `using` CUDA.jl and CUDSS.jl. Then `psolver_direct` will automatically use
+ the CUDSS solver.
diff --git a/docs/src/features/les.md b/docs/src/features/les.md
index 952b22cc0..9b1e741ac 100644
--- a/docs/src/features/les.md
+++ b/docs/src/features/les.md
@@ -1,3 +1,7 @@
+```@meta
+CurrentModule = IncompressibleNavierStokes
+```
+
# Large eddy simulation
Depending on the problem specification, a given grid resolution may not be
@@ -28,12 +32,7 @@ atomic scales. The new turbulent viscosity on the other hand, models energy
transfer from resolved to unresolved scales. This non-constant field is
computed from the local velocity field.
-The following eddy viscosity models are available:
-
-- [`SmagorinskyModel`](@ref)
-- [`QRModel`](@ref)
-- [`MixingLengthModel`](@ref)
-
-In addition, the default [`LaminarModel`](@ref) assumes that there are no
-sub-grid stresses. It can be used if the grid is sufficiently refined for the
-given flow. It has the advantage of having a constant diffusion operator.
+```@docs
+smagtensor!
+smagorinsky_closure
+```
diff --git a/docs/src/features/operators.md b/docs/src/features/operators.md
new file mode 100644
index 000000000..54ac99408
--- /dev/null
+++ b/docs/src/features/operators.md
@@ -0,0 +1,49 @@
+```@meta
+CurrentModule = IncompressibleNavierStokes
+```
+
+# Operators
+
+All discrete operators are built using
+[KernelAbstractions.jl](https://github.com/JuliaGPU/KernelAbstractions.jl/)
+and Cartesian indices, similar to
+[WaterLily.jl](https://github.com/weymouth/WaterLily.jl/).
+This allows for dimension- and backend-agnostic code. See this
+[blog post](https://b-fg.github.io/2023/05/07/waterlily-on-gpu.html)
+for how to write kernels. IncompressibleNavierStokes previously relied on
+assembling sparse operators to perform the same operations. While being very
+efficient and also compatible with CUDA (CUSPARSE), storing these matrices in
+memory is expensive for large 3D problems.
+
+```@docs
+Offset
+divergence!
+divergence
+vorticity
+vorticity!
+convection!
+diffusion!
+bodyforce!
+momentum!
+momentum
+laplacian!
+laplacian
+pressuregradient!
+pressuregradient
+interpolate_u_p
+interpolate_u_p!
+interpolate_ω_p
+interpolate_ω_p!
+Dfield!
+Dfield
+Qfield!
+Qfield
+eig2field!
+eig2field
+kinetic_energy
+kinetic_energy!
+total_kinetic_energy
+tensorbasis
+divoftensor!
+tensorbasis!
+```
diff --git a/docs/src/features/precision.md b/docs/src/features/precision.md
index dd794ae08..f880f0ff2 100644
--- a/docs/src/features/precision.md
+++ b/docs/src/features/precision.md
@@ -19,6 +19,6 @@ an example on how to enforce floating point type hygiene.
!!! note "Pressure solvers"
[`SparseArrays.jl`](https://github.com/JuliaSparse/SparseArrays.jl)s
sparse matrix factorizations only support double precision.
- [`DirectPressureSolver`](@ref) only works for `Float64`. Consider
- using an iterative solver such as [`CGPressureSolver`](@ref)
+ [`psolver_direct`](@ref) only works for `Float64`. Consider
+ using an iterative solver such as [`psolver_cg`](@ref)
when using single or half precision.
diff --git a/docs/src/features/pressure.md b/docs/src/features/pressure.md
index 80453f5d9..a509c4bdb 100644
--- a/docs/src/features/pressure.md
+++ b/docs/src/features/pressure.md
@@ -1,15 +1,28 @@
+```@meta
+CurrentModule = IncompressibleNavierStokes
+```
+
# Pressure solvers
The discrete pressure Poisson equation
```math
-A p_h = g
+L p = W M F(u)
```
-enforces divergence freeness. There are three options for solving this system:
+enforces divergence freeness. There are multiple options for solving this
+system.
-- [`DirectPressureSolver`](@ref) factorizes the Laplace matrix ``A`` such that
- the system can be solved for different right hand sides. This currently
- only works for double precision on the CPU.
-- [`CGPressureSolver`](@ref) uses conjugate gradients to solve the system for
- different ``f``.
-- [`SpectralPressureSolver`](@ref) solves the system in Fourier space, but only
- in the case of a uniform grid with periodic boundary conditions.
+```@docs
+default_psolver
+psolver_direct
+psolver_cg
+psolver_cg_matrix
+psolver_spectral
+psolver_spectral_lowmemory
+pressure
+pressure!
+poisson
+poisson!
+applypressure!
+project
+project!
+```
diff --git a/docs/src/features/steppers.md b/docs/src/features/steppers.md
deleted file mode 100644
index d9d322fc4..000000000
--- a/docs/src/features/steppers.md
+++ /dev/null
@@ -1,9 +0,0 @@
-# Time steppers
-
-IncompressibleNavierStokes provides a collection of explicit and implicit
-[Runge-Kutta methods](../api/tableaux.md), in addition to Adams-Bashforth
-Crank-Nicolson and one-leg beta method time steppers.
-
-The code is currently not adapted to time steppers from
-[DifferentialEquations.jl](https://docs.sciml.ai/DiffEqDocs/stable/solvers/dae_solve/),
-but they may be integrated in the future.
diff --git a/docs/src/features/temperature.md b/docs/src/features/temperature.md
new file mode 100644
index 000000000..5bfe77ce7
--- /dev/null
+++ b/docs/src/features/temperature.md
@@ -0,0 +1,33 @@
+```@meta
+CurrentModule = IncompressibleNavierStokes
+```
+
+# Temperature equation
+
+IncompressibleNavierStokes.jl supports adding a temperature equation, which is
+coupled back to the momentum equation through a gravity term
+[Sanderse2023](@cite).
+
+To enable the temperature equation, you need to set the `temperature` keyword
+in setup:
+
+```julia
+setup = Setup(
+ args...;
+ kwargs...,
+ temperature = temperature_equation(; kwargs...),
+)
+```
+
+where `temperature_equation` can be configured as follows:
+
+```@docs
+temperature_equation
+```
+
+Some operators are available for the temperature equation:
+```@docs
+gravity!
+convection_diffusion_temp!
+dissipation!
+```
diff --git a/docs/src/getting_started.md b/docs/src/getting_started.md
index eac7cc05d..d005edf7a 100644
--- a/docs/src/getting_started.md
+++ b/docs/src/getting_started.md
@@ -8,4 +8,4 @@ add IncompressibleNavierStokes
```
which will install the package and all dependencies to your local environment.
-Note that IncompressibleNavierStokes requires Julia version `1.7` or above.
+Note that IncompressibleNavierStokes requires Julia version `1.9` or above.
diff --git a/examples/Actuator2D.jl b/examples/Actuator2D.jl
index ca3434f88..a1470b914 100644
--- a/examples/Actuator2D.jl
+++ b/examples/Actuator2D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Unsteady actuator case - 2D
#
# In this example, an unsteady inlet velocity profile at encounters a wind
@@ -21,94 +14,73 @@ end #src
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "Actuator2D"
-
-# Viscosity model
-viscosity_model = LaminarModel(; Re = 100.0)
-
-# Boundary conditions: Unsteady BC requires time derivatives
-u_bc(x, y, t) = x ≈ 0.0 ? cos(π / 6 * sin(π / 6 * t)) : 0.0
-v_bc(x, y, t) = x ≈ 0.0 ? sin(π / 6 * sin(π / 6 * t)) : 0.0
-dudt_bc(x, y, t) = x ≈ 0.0 ? -(π / 6)^2 * cos(π / 6 * t) * sin(π / 6 * sin(π / 6 * t)) : 0.0
-dvdt_bc(x, y, t) = x ≈ 0.0 ? (π / 6)^2 * cos(π / 6 * t) * cos(π / 6 * sin(π / 6 * t)) : 0.0
-bc_type = (;
- u = (; x = (:dirichlet, :pressure), y = (:symmetric, :symmetric)),
- v = (; x = (:dirichlet, :symmetric), y = (:pressure, :pressure)),
-)
+# Output directory
+output = "output/Actuator2D"
# A 2D grid is a Cartesian product of two vectors
n = 40
-x = LinRange(0.0, 10.0, 5n)
-y = LinRange(-2.0, 2.0, 2n)
-plot_grid(x, y)
+x = LinRange(0.0, 10.0, 5n + 1)
+y = LinRange(-2.0, 2.0, 2n + 1)
+plotgrid(x, y)
+
+# Boundary conditions
+boundary_conditions = (
+ ## x left, x right
+ (
+ ## Unsteady BC requires time derivatives
+ DirichletBC(
+ (dim, x, y, t) -> sin(π / 6 * sin(π / 6 * t) + π / 2 * (dim() == 1)),
+ (dim, x, y, t) ->
+ (π / 6)^2 *
+ cos(π / 6 * t) *
+ cos(π / 6 * sin(π / 6 * t) + π / 2 * (dim() == 1)),
+ ),
+ PressureBC(),
+ ),
+
+ ## y rear, y front
+ (PressureBC(), PressureBC()),
+)
# Actuator body force: A thrust coefficient `Cₜ` distributed over a thin rectangle
xc, yc = 2.0, 0.0 # Disk center
D = 1.0 # Disk diameter
δ = 0.11 # Disk thickness
-Cₜ = 5e-4 # Thrust coefficient
+Cₜ = 0.2 # Thrust coefficient
cₜ = Cₜ / (D * δ)
inside(x, y) = abs(x - xc) ≤ δ / 2 && abs(y - yc) ≤ D / 2
-bodyforce_u(x, y) = -cₜ * inside(x, y)
-bodyforce_v(x, y) = 0.0
+bodyforce(dim, x, y, t) = dim() == 1 ? -cₜ * inside(x, y) : 0.0
# Build setup and assemble operators
-setup = Setup(
- x,
- y;
- viscosity_model,
- u_bc,
- v_bc,
- dudt_bc,
- dvdt_bc,
- bc_type,
- bodyforce_u,
- bodyforce_v,
-);
-
-# Time interval
-t_start, t_end = tlims = (0.0, 12.0)
+setup = Setup(x, y; Re = 100.0, boundary_conditions, bodyforce);
# Initial conditions (extend inflow)
-initial_velocity_u(x, y) = 1.0
-initial_velocity_v(x, y) = 0.0
-initial_pressure(x, y) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t_start;
- initial_pressure,
-);
-
-# Iteration processors
-processors = (
- field_plotter(setup; nupdate = 1),
- ## energy_history_plotter(setup; nupdate = 10),
- ## energy_spectrum_plotter(setup; nupdate = 10),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
-);
+ustart = create_initial_conditions(setup, (dim, x, y) -> dim() == 1 ? 1.0 : 0.0);
# Solve unsteady problem
-V, p, outputs = solve_unsteady(
+state, outputs = solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- method = RK44P2(),
+ ustart,
+ tlims = (0.0, 12.0),
+ method = RKMethods.RK44P2(),
Δt = 0.05,
- processors,
- inplace = true,
+ processors = (
+ rtp = realtimeplotter(; setup, plot = fieldplot, nupdate = 1),
+ ## ehist = realtimeplotter(; setup, plot = energy_history_plot, nupdate = 1),
+ ## espec = realtimeplotter(; setup, plot = energy_spectrum_plot, nupdate = 1),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = "$output", filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
);
-#md current_figure()
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`.
+# We may visualize or export the computed fields `(u, p)`.
+
+# Export to VTK
+save_vtk(setup, state.u, state.t, "$output/solution")
# We create a box to visualize the actuator.
box = (
@@ -116,30 +88,17 @@ box = (
[yc + D / 2, yc - D / 2, yc - D / 2, yc + D / 2, yc + D / 2],
)
-# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
-
# Plot pressure
-fig = plot_pressure(setup, p)
+fig = fieldplot(state; setup, fieldname = :pressure)
lines!(box...; color = :red)
fig
# Plot velocity
-fig = plot_velocity(setup, V, t_end)
+fig = fieldplot(state; setup, fieldname = :velocity)
lines!(box...; color = :red)
fig
# Plot vorticity
-fig = plot_vorticity(setup, V, t_end)
-lines!(box...; color = :red)
-fig
-
-# Plot streamfunction
-fig = plot_streamfunction(setup, V, t_end)
-lines!(box...; color = :red)
-fig
-
-# Plot force
-fig = plot_force(setup, t_end)
+fig = fieldplot(state; setup, fieldname = :vorticity)
lines!(box...; color = :red)
fig
diff --git a/examples/Actuator3D.jl b/examples/Actuator3D.jl
index 9966d53c9..01792dd38 100644
--- a/examples/Actuator3D.jl
+++ b/examples/Actuator3D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Unsteady actuator case - 3D
#
# In this example, an unsteady inlet velocity profile at encounters a wind
@@ -21,134 +14,94 @@ end #src
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "Actuator3D"
+# Output directory
+output = "output/Actuator3D"
+
+# Floating point type
+T = Float64
+
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
-# Viscosity model
-viscosity_model = LaminarModel(; Re = 100.0)
+# Reynolds number
+Re = T(100)
+
+# A 3D grid is a Cartesian product of three vectors
+x = LinRange(0.0, 6.0, 31)
+y = LinRange(-2.0, 2.0, 41)
+z = LinRange(-2.0, 2.0, 41)
+plotgrid(x, y, z)
# Boundary conditions: Unsteady BC requires time derivatives
-u_bc(x, y, z, t) = x ≈ 0.0 ? cos(π / 6 * sin(π / 6 * t)) : 0.0
-v_bc(x, y, z, t) = x ≈ 0.0 ? sin(π / 6 * sin(π / 6 * t)) : 0.0
-w_bc(x, y, z, t) = 0.0
-dudt_bc(x, y, z, t) =
- x ≈ 0.0 ? -(π / 6)^2 * cos(π / 6 * t) * sin(π / 6 * sin(π / 6 * t)) : 0.0
-dvdt_bc(x, y, z, t) =
- x ≈ 0.0 ? (π / 6)^2 * cos(π / 6 * t) * cos(π / 6 * sin(π / 6 * t)) : 0.0
-dwdt_bc(x, y, z, t) = 0.0
-bc_type = (;
- u = (;
- x = (:dirichlet, :pressure),
- y = (:symmetric, :symmetric),
- z = (:symmetric, :symmetric),
- ),
- v = (;
- x = (:dirichlet, :symmetric),
- y = (:pressure, :pressure),
- z = (:symmetric, :symmetric),
+boundary_conditions = (
+ ## x left, x right
+ (
+ DirichletBC(
+ (dim, x, y, z, t) ->
+ dim() == 1 ? cos(π / 6 * sin(π / 6 * t)) :
+ dim() == 2 ? sin(π / 6 * sin(π / 6 * t)) : zero(x),
+ (dim, x, y, z, t) ->
+ dim() == 1 ? -(π / 6)^2 * cos(π / 6 * t) * sin(π / 6 * sin(π / 6 * t)) :
+ dim() == 2 ? (π / 6)^2 * cos(π / 6 * t) * cos(π / 6 * sin(π / 6 * t)) :
+ zero(x),
+ ),
+ PressureBC(),
),
- w = (;
- x = (:dirichlet, :symmetric),
- y = (:symmetric, :symmetric),
- z = (:pressure, :pressure),
- ),
-)
-# A 3D grid is a Cartesian product of three vectors
-x = LinRange(0.0, 6.0, 30)
-y = LinRange(-2.0, 2.0, 40)
-z = LinRange(-2.0, 2.0, 40)
-plot_grid(x, y, z)
+ ## y rear, y front
+ (PressureBC(), PressureBC()),
+
+ ## z rear, z front
+ (PressureBC(), PressureBC()),
+)
# Actuator body force: A thrust coefficient `Cₜ` distributed over a short cylinder
-cx, cy, cz = 2.0, 0.0, 0.0 # Disk center
-D = 1.0 # Disk diameter
-δ = 0.11 # Cylinder height
-Cₜ = 5e-4 # Thrust coefficient
+cx, cy, cz = T(2), T(0), T(0) # Disk center
+D = T(1) # Disk diameter
+δ = T(0.11) # Disk thickness
+Cₜ = T(0.2) # Thrust coefficient
cₜ = Cₜ / (π * (D / 2)^2 * δ)
inside(x, y, z) = abs(x - cx) ≤ δ / 2 && (y - cy)^2 + (z - cz)^2 ≤ (D / 2)^2
-bodyforce_u(x, y, z) = -cₜ * inside(x, y, z)
-bodyforce_v(x, y, z) = 0.0
-bodyforce_w(x, y, z) = 0.0
+bodyforce(dim, x, y, z) = dim() == 1 ? -cₜ * inside(x, y, z) : zero(x)
# Build setup and assemble operators
-setup = Setup(
- x,
- y,
- z;
- viscosity_model,
- u_bc,
- v_bc,
- w_bc,
- dudt_bc,
- dvdt_bc,
- dwdt_bc,
- bc_type,
- bodyforce_u,
- bodyforce_v,
- bodyforce_w,
-);
-
-# Time interval
-t_start, t_end = tlims = (0.0, 3.0)
+setup = Setup(x, y, z; Re, boundary_conditions, bodyforce, ArrayType);
# Initial conditions (extend inflow)
-initial_velocity_u(x, y, z) = 1.0
-initial_velocity_v(x, y, z) = 0.0
-initial_velocity_w(x, y, z) = 0.0
-initial_pressure(x, y, z) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- initial_velocity_w,
- t_start;
- initial_pressure,
-);
-
-# Iteration processors
-processors = (
- field_plotter(setup; nupdate = 5),
- ## energy_history_plotter(setup; nupdate = 10),
- ## energy_spectrum_plotter(setup; nupdate = 10),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 2, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
-);
+ustart = create_initial_conditions(setup, (dim, x, y, z) -> dim() == 1 ? one(x) : zero(x));
# Solve unsteady problem
-V, p, outputs = solve_unsteady(
+(; u, t), outputs = solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- method = RK44P2(),
- Δt = 0.05,
- processors,
- inplace = true,
+ ustart,
+ tlims = (T(0), T(3)),
+ method = RKMethods.RK44P2(),
+ Δt = T(0.05),
+ processors = (
+ rtp = realtimeplotter(;
+ setup,
+ plot = fieldplot,
+ ## plot = energy_history_plot,
+ ## plot = energy_spectrum_plot,
+ nupdate = 1,
+ ),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = "$output", filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
);
-#md current_figure()
# ## Post-process
#
# We may visualize or export the computed fields `(V, p)`
-# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
-
-# Plot pressure
-plot_pressure(setup, p)
-
-# Plot velocity
-plot_velocity(setup, V, t_end)
+# Field plot
+outputs.rtp
-# Plot vorticity
-plot_vorticity(setup, V, t_end)
-
-# Plot streamfunction
-## plot_streamfunction(setup, V, t_end)
-nothing
-
-# Plot force
-plot_force(setup, t_end)
+# Export to VTK
+save_vtk(setup, u, t, "$output/solution")
diff --git a/examples/BackwardFacingStep2D.jl b/examples/BackwardFacingStep2D.jl
index fe539ecad..a93023193 100644
--- a/examples/BackwardFacingStep2D.jl
+++ b/examples/BackwardFacingStep2D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Backward Facing Step - 2D
#
# In this example we consider a channel with walls at the top and bottom, and a
@@ -22,77 +15,83 @@ end #src
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "BackwardFacingStep2D"
+# Output directory
+output = "output/BackwardFacingStep2D"
+
+# Floating point type
+T = Float64
+
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
-# Viscosity model
-viscosity_model = LaminarModel(; Re = 3000.0)
+# Reynolds number
+Re = T(3_000)
# Boundary conditions: steady inflow on the top half
-u_bc(x, y, t) = x ≈ 0 && y ≥ 0 ? 24y * (1 / 2 - y) : 0.0
-v_bc(x, y, t) = 0.0
-bc_type = (;
- u = (; x = (:dirichlet, :pressure), y = (:dirichlet, :dirichlet)),
- v = (; x = (:dirichlet, :symmetric), y = (:dirichlet, :dirichlet)),
+U(dim, x, y, t) =
+ dim() == 1 && y ≥ 0 ? 24y * (one(x) / 2 - y) : zero(x) + randn(typeof(x)) / 1_000
+dUdt(dim, x, y, t) = zero(x)
+boundary_conditions = (
+ ## x left, x right
+ (DirichletBC(U, dUdt), PressureBC()),
+
+ ## y rear, y front
+ (DirichletBC(), DirichletBC()),
)
# A 2D grid is a Cartesian product of two vectors. Here we refine the grid near
# the walls.
-x = LinRange(0.0, 10.0, 300)
-y = cosine_grid(-0.5, 0.5, 50)
-plot_grid(x, y)
+x = LinRange(T(0), T(10), 301)
+y = cosine_grid(-T(0.5), T(0.5), 51)
+plotgrid(x, y)
# Build setup and assemble operators
-setup = Setup(x, y; viscosity_model, u_bc, v_bc, bc_type);
-
-# Time interval
-t_start, t_end = tlims = (0.0, 7.0)
+setup = Setup(x, y; Re, boundary_conditions, ArrayType);
# Initial conditions (extend inflow)
-initial_velocity_u(x, y) = y ≥ 0.0 ? 24y * (1 / 2 - y) : 0.0
-initial_velocity_v(x, y) = 0.0
-initial_pressure(x, y) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t_start;
- initial_pressure,
-);
+ustart = create_initial_conditions(setup, (dim, x, y) -> U(dim, x, y, zero(x)));
# Solve steady state problem
-V, p = solve_steady_state(setup, V₀, p₀);
-
-# Iteration processors
-processors = (
- field_plotter(setup; nupdate = 5),
- ## energy_history_plotter(setup; nupdate = 10),
- ## energy_spectrum_plotter(setup; nupdate = 10),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 20, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
-);
+## u, p = solve_steady_state(setup, u₀, p₀);
+nothing
# Solve unsteady problem
-V, p, outputs = solve_unsteady(setup, V₀, p₀, tlims; Δt = 0.002, processors, inplace = true)
-#md current_figure()
+state, outputs = solve_unsteady(;
+ setup,
+ ustart,
+ tlims = (T(0), T(7)),
+ Δt = T(0.002),
+ processors = (
+ rtp = realtimeplotter(;
+ setup,
+ plot = fieldplot,
+ ## plot = energy_history_plot,
+ ## plot = energy_spectrum_plot,
+ nupdate = 1,
+ ),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
+);
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`
+# We may visualize or export the computed fields
# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
+save_vtk(setup, state.u, state.t, "$output/solution")
# Plot pressure
-plot_pressure(setup, p)
+fieldplot(state; setup, fieldname = :pressure)
# Plot velocity
-plot_velocity(setup, V, t_end)
+fieldplot(state; setup, fieldname = :velocity)
# Plot vorticity
-plot_vorticity(setup, V, t_end)
-
-# Plot streamfunction
-plot_streamfunction(setup, V, t_end)
+fieldplot(state; setup, fieldname = :vorticity)
diff --git a/examples/BackwardFacingStep3D.jl b/examples/BackwardFacingStep3D.jl
index 861c9f78d..72ba438d8 100644
--- a/examples/BackwardFacingStep3D.jl
+++ b/examples/BackwardFacingStep3D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Backward Facing Step - 3D
#
# In this example we consider a channel with periodic side boundaries, walls at
@@ -22,94 +15,85 @@ end #src
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "BackwardFacingStep3D"
+# Output directory
+output = "output/BackwardFacingStep3D"
-# Viscosity model
-viscosity_model = LaminarModel(; Re = 3000.0)
+# Floating point type
+T = Float64
-# Boundary conditions: steady inflow on the top half
-u_bc(x, y, z, t) = x ≈ 0 && y ≥ 0 ? 24y * (1 / 2 - y) : 0.0
-v_bc(x, y, z, t) = 0.0
-w_bc(x, y, z, t) = 0.0
-bc_type = (;
- u = (;
- x = (:dirichlet, :pressure),
- y = (:dirichlet, :dirichlet),
- z = (:periodic, :periodic),
- ),
- v = (;
- x = (:dirichlet, :symmetric),
- y = (:dirichlet, :dirichlet),
- z = (:periodic, :periodic),
- ),
- w = (;
- x = (:dirichlet, :symmetric),
- y = (:dirichlet, :dirichlet),
- z = (:periodic, :periodic),
- ),
-)
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+# Reynolds number
+Re = T(3000)
# A 3D grid is a Cartesian product of three vectors
-x = LinRange(0, 10, 160)
-y = LinRange(-0.5, 0.5, 16)
-z = LinRange(-0.25, 0.25, 8)
-plot_grid(x, y, z)
+x = LinRange(T(0), T(10), 129)
+y = LinRange(-T(0.5), T(0.5), 17)
+z = LinRange(-T(0.25), T(0.25), 9)
+plotgrid(x, y, z)
-# Build setup and assemble operators
-setup = Setup(x, y, z; viscosity_model, u_bc, v_bc, w_bc, bc_type);
+# Boundary conditions: steady inflow on the top half
+U(dim, x, y, z, t) = dim() == 1 && y ≥ 0 ? 24y * (one(x) / 2 - y) : zero(x)
+dUdt(dim, x, y, z, t) = zero(x)
+boundary_conditions = (
+ ## x left, x right
+ (DirichletBC(U, dUdt), PressureBC()),
+
+ ## y rear, y front
+ (DirichletBC(), DirichletBC()),
-# Time interval
-t_start, t_end = tlims = (0.0, 7.0)
+ ## z bottom, z top
+ (PeriodicBC(), PeriodicBC()),
+)
+
+# Build setup and assemble operators
+setup = Setup(x, y, z; Re, boundary_conditions, ArrayType);
# Initial conditions (extend inflow)
-initial_velocity_u(x, y, z) = y ≥ 0 ? 24y * (1 / 2 - y) : 0.0
-initial_velocity_v(x, y, z) = 0.0
-initial_velocity_w(x, y, z) = 0.0
-initial_pressure(x, y, z) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- initial_velocity_w,
- t_start;
- initial_pressure,
-);
+ustart = create_initial_conditions(setup, (dim, x, y, z) -> U(dim, x, y, z, zero(x)));
# Solve steady state problem
-V, p = solve_steady_state(setup, V₀, p₀);
-
-# Iteration processors
-processors = (
- field_plotter(setup; nupdate = 50),
- ## energy_history_plotter(setup; nupdate = 10),
- ## energy_spectrum_plotter(setup; nupdate = 10),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 20, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 10),
-);
+## u, p = solve_steady_state(setup, u₀, p₀);
+nothing
# Solve unsteady problem
-V, p, outputs = solve_unsteady(setup, V₀, p₀, tlims; Δt = 0.01, processors, inplace = true)
-#md current_figure()
+state, outputs = solve_unsteady(;
+ setup,
+ ustart,
+ tlims = (T(0), T(7)),
+ Δt = T(0.01),
+ processors = (
+ rtp = realtimeplotter(;
+ setup,
+ plot = fieldplot,
+ ## plot = energy_history_plot,
+ ## plot = energy_spectrum_plot,
+ nupdate = 1,
+ ),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
+)
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`
+# We may visualize or export the computed fields
# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
+save_vtk(setup, state.u, state.t, "$output/solution")
# Plot pressure
-plot_pressure(setup, p)
+fieldplot(state; setup, fieldname = :pressure)
# Plot velocity
-plot_velocity(setup, V, t_end)
+fieldplot(state; setup, fieldname = :velocity)
# Plot vorticity
-plot_vorticity(setup, V, t_end)
-
-# Plot streamfunction
-## plot_streamfunction(setup, V, t_end)
-nothing
+fieldplot(state; setup, fieldname = :vorticity)
diff --git a/examples/DecayingTurbulence2D.jl b/examples/DecayingTurbulence2D.jl
index ade9faa93..e3d5beb9f 100644
--- a/examples/DecayingTurbulence2D.jl
+++ b/examples/DecayingTurbulence2D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Decaying Homogeneous Isotropic Turbulence - 2D
#
# In this example we consider decaying homogeneous isotropic turbulence,
@@ -13,99 +6,71 @@ end #src
# specific energy spectrum. Due to viscous dissipation, the turbulent features
# eventually group to form larger visible eddies.
-# We start by loading packages.
-# A [Makie](https://github.com/JuliaPlots/Makie.jl) plotting backend is needed
-# for plotting. `GLMakie` creates an interactive window (useful for real-time
-# plotting), but does not work when building this example on GitHub.
-# `CairoMakie` makes high-quality static vector-graphics plots.
-
#md using CairoMakie
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "DecayingTurbulence2D"
+# Output directory
+output = "output/DecayingTurbulence2D"
# Floating point precision
-T = Float32
-
-# To use CPU: Do not move any arrays
-device = identity
+T = Float64
-# To use GPU, use `cu` to move arrays to the GPU.
-# Note: `cu` converts to Float32
-## using CUDA
-## device = cu
-nothing
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
# Viscosity model
-viscosity_model = LaminarModel(; Re = T(10_000))
+Re = T(10_000)
# A 2D grid is a Cartesian product of two vectors
n = 256
-lims = (T(0), T(1))
-x = LinRange(lims..., n + 1)
-y = LinRange(lims..., n + 1)
-# plot_grid(x, y)
+lims = T(0), T(1)
+x = LinRange(lims..., n + 1), LinRange(lims..., n + 1)
# Build setup and assemble operators
-setup = Setup(x, y; viscosity_model);
-
-# Since the grid is uniform and identical for x and y, we may use a specialized
-# spectral pressure solver
-pressure_solver = SpectralPressureSolver(setup);
-
-# Initial conditions
-V₀, p₀ = random_field(setup; A = T(1_000_000), σ = T(30), s = 5, pressure_solver);
-
-# Iteration processors
-processors = (
- field_plotter(device(setup); nupdate = 20),
- energy_history_plotter(device(setup); nupdate = 20, displayfig = false),
- energy_spectrum_plotter(device(setup); nupdate = 20, displayfig = false),
- ## animator(device(setup), "vorticity.mp4"; nupdate = 16),
- ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 100),
-);
+setup = Setup(x...; Re, ArrayType);
-# Time interval
-t_start, t_end = tlims = (T(0), T(1.0))
+# Create random initial conditions
+ustart = random_field(setup, T(0));
# Solve unsteady problem
-V, p, outputs = solve_unsteady(
+state, outputs = solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- Δt = T(0.001),
- processors,
- pressure_solver,
- inplace = true,
- device,
+ ustart,
+ tlims = (T(0), T(1)),
+ Δt = T(1e-3),
+ processors = (
+ ## rtp = realtimeplotter(; setup, nupdate = 1),
+ ehist = realtimeplotter(;
+ setup,
+ plot = energy_history_plot,
+ nupdate = 10,
+ displayfig = false,
+ ),
+ espec = realtimeplotter(; setup, plot = energy_spectrum_plot, nupdate = 10),
+ ## anim = animator(; setup, path = "$output/solution.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 100),
+ ),
);
-# Field plot
-outputs[1]
-
-# Energy history plot
-outputs[2]
-
-# Energy spectrum plot
-outputs[3]
-
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`
+# We may visualize or export the computed fields `(u, p)`
-# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
+# Energy history
+outputs.ehist
-# Plot pressure
-plot_pressure(setup, p)
+# Energy spectrum
+outputs.espec
-# Plot velocity
-plot_velocity(setup, V, t_end)
+# Export to VTK
+save_vtk(setup, state.u, state.t, "$output/solution")
-# Plot vorticity
-plot_vorticity(setup, V, t_end)
+# Plot field
+fieldplot(state; setup)
diff --git a/examples/DecayingTurbulence3D.jl b/examples/DecayingTurbulence3D.jl
index 33897d2ed..684177111 100644
--- a/examples/DecayingTurbulence3D.jl
+++ b/examples/DecayingTurbulence3D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Decaying Homogeneous Isotropic Turbulence - 3D
#
# In this example we consider decaying homogeneous isotropic turbulence,
@@ -13,100 +6,75 @@ end #src
# specific energy spectrum. Due to viscous dissipation, the turbulent features
# eventually group to form larger visible eddies.
-# We start by loading packages.
-# A [Makie](https://github.com/JuliaPlots/Makie.jl) plotting backend is needed for plotting. `GLMakie` creates an interactive window (useful for real-time
-# plotting), but does not work when building this example on GitHub.
-# `CairoMakie` makes high-quality static vector-graphics plots.
-
-using FFTW
#md using CairoMakie
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "DecayingTurbulence3D"
+# Output directory
+output = "output/DecayingTurbulence3D"
# Floating point precision
-T = Float32
+T = Float64
-# To use CPU: Do not move any arrays
-device = identity
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
-# To use GPU, use `cu` to move arrays to the GPU.
-# Note: `cu` converts to Float32
-## using CUDA
-## device = cu
-nothing
-
-# Viscosity model
-viscosity_model = LaminarModel(; Re = T(10_000))
+# Reynolds number
+Re = T(6_000)
# A 3D grid is a Cartesian product of three vectors
n = 32
-lims = (T(0), T(1))
+lims = T(0), T(1)
x = LinRange(lims..., n + 1)
y = LinRange(lims..., n + 1)
z = LinRange(lims..., n + 1)
-# plot_grid(x, y, z)
# Build setup and assemble operators
-setup = Setup(x, y, z; viscosity_model);
+setup = Setup(x, y, z; Re, ArrayType);
# Since the grid is uniform and identical for x, y, and z, we may use a
# specialized spectral pressure solver
-pressure_solver = SpectralPressureSolver(setup);
+psolver = psolver_spectral(setup);
# Initial conditions
-V₀, p₀ = random_field(setup; A = T(1_000_000), σ = T(30), s = 5, pressure_solver)
-
-# Time interval
-t_start, t_end = tlims = (T(0), T(1.0))
-
-# Iteration processors
-processors = (
- field_plotter(device(setup); nupdate = 10),
- energy_history_plotter(device(setup); nupdate = 10),
- energy_spectrum_plotter(device(setup); nupdate = 10),
- ## animator(device(setup), "vorticity.mp4"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
-);
+ustart = random_field(setup; psolver);
# Solve unsteady problem
-V, p, outputs = solve_unsteady(
+(; u, t), outputs = solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- Δt = T(0.001),
- processors,
- pressure_solver,
- inplace = true,
- device,
+ ustart,
+ tlims = (T(0), T(1)),
+ Δt = T(1e-3),
+ psolver,
+ processors = (
+ ## rtp = realtimeplotter(; setup, plot = fieldplot, nupdate = 10),
+ ehist = realtimeplotter(;
+ setup,
+ plot = energy_history_plot,
+ nupdate = 10,
+ displayfig = false,
+ ),
+ espec = realtimeplotter(; setup, plot = energy_spectrum_plot, nupdate = 10),
+ ## anim = animator(; setup, path = "$output/solution.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 100),
+ ),
);
-# Field plot
-outputs[1]
+# ## Post-process
+#
+# We may visualize or export the computed fields `(u, p)`
# Energy history
-outputs[2]
+outputs.ehist
# Energy spectrum
-outputs[3]
-
-# ## Post-process
-#
-# We may visualize or export the computed fields `(V, p)`
+outputs.espec
# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
-
-# Plot pressure
-plot_pressure(setup, p)
-
-# Plot velocity
-plot_velocity(setup, V, t_end)
-
-# Plot vorticity
-plot_vorticity(setup, V, t_end)
+save_vtk(setup, u, t, "$output/solution")
diff --git a/examples/LidDrivenCavity2D.jl b/examples/LidDrivenCavity2D.jl
index d1a7550ef..e020b7edf 100644
--- a/examples/LidDrivenCavity2D.jl
+++ b/examples/LidDrivenCavity2D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Tutorial: Lid-Driven Cavity - 2D
#
# In this example we consider a box with a moving lid. The velocity is
@@ -14,6 +7,7 @@ end #src
# We start by loading packages.
# A [Makie](https://github.com/JuliaPlots/Makie.jl) plotting backend is needed
+#
# for plotting. `GLMakie` creates an interactive window (useful for real-time
# plotting), but does not work when building this example on GitHub.
# `CairoMakie` makes high-quality static vector-graphics plots.
@@ -23,7 +17,7 @@ using GLMakie #!md
using IncompressibleNavierStokes
# Case name for saving results
-name = "LidDrivenCavity2D"
+output = "output/LidDrivenCavity2D"
# The code allows for using different floating point number types, including single
# precision (`Float32`) and double precision (`Float64`). On the CPU, the speed
@@ -33,86 +27,60 @@ name = "LidDrivenCavity2D"
# scaled judiciously to avoid vanishing digits when applying differential
# operators of the form "right minus left divided by small distance".
-## T = Float16
-## T = Float32
-T = Float64
-
# Note how floating point type hygiene is enforced in the following using `T`
# to avoid mixing different precisions.
-# We can also choose to do the computations on a different device.
-# By default, the computations are performed on the host (CPU). An optional
-# `device` keyword allows for moving arrays and operator to a different device
-# such as a GPU. Currently, only Nvidia GPUs with CUDA support sparse operators
-# and fast Fourier transform used by IncompressibleNavierStokes.
-
-# For CPU
-device = identity
-
-# For GPU (note that `cu` converts to `Float32`)
-## using CUDA
-## device = cu
-nothing
+# T = Float64
+T = Float32
+## T = Float16
-# Available viscosity models are:
-#
-# - [`LaminarModel`](@ref),
-# - [`MixingLengthModel`](@ref),
-# - [`SmagorinskyModel`](@ref), and
-# - [`QRModel`](@ref).
+# We can also choose to do the computations on a different device. By default,
+# the computations are performed on the host (CPU). An optional `ArrayType`
+# allows for moving arrays to a different device such as a GPU.
#
-# They all take a Reynolds number as a parameter. Here we choose a moderate
-# Reynolds number. Note how we pass the floating point type.
-viscosity_model = LaminarModel(; Re = T(1_000))
-
-# Dirichlet boundary conditions are specified as plain Julia functions. They
-# are marked by the `:dirichlet` symbol. Other possible BC types are
-# `:periodic`, `:symmetric`, and `:pressure`.
-u_bc(x, y, t) = y ≈ 1 ? T(1) : T(0)
-v_bc(x, y, t) = T(0)
-bc_type = (;
- u = (; x = (:dirichlet, :dirichlet), y = (:dirichlet, :dirichlet)),
- v = (; x = (:dirichlet, :dirichlet), y = (:dirichlet, :dirichlet)),
+# Note: For GPUs, single precision is preferred.
+
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+# Here we choose a moderate Reynolds number. Note how we pass the floating point type.
+Re = T(1_000)
+
+# Non-zero Dirichlet boundary conditions are specified as plain Julia functions.
+# Note that time derivatives are required.
+boundary_conditions = (
+ ## x left, x right
+ (DirichletBC(), DirichletBC()),
+
+ ## y bottom, y top
+ (
+ DirichletBC(),
+ DirichletBC(
+ (dim, x, y, t) -> dim() == 1 ? one(x) : zero(x),
+ (dim, x, y, t) -> zero(x),
+ ),
+ ),
)
# We create a two-dimensional domain with a box of size `[1, 1]`. The grid is
# created as a Cartesian product between two vectors. We add a refinement near
# the walls.
-n = 40
-lims = (T(0), T(1))
+n = 32
+lims = T(0), T(1)
x = cosine_grid(lims..., n)
y = cosine_grid(lims..., n)
-plot_grid(x, y)
+plotgrid(x, y)
# We can now build the setup and assemble operators.
# A 3D setup is built if we also provide a vector of z-coordinates.
-setup = Setup(x, y; viscosity_model, u_bc, v_bc, bc_type);
-
-# The pressure solver is used to solve the pressure Poisson equation.
-# Available solvers are
-#
-# - [`DirectPressureSolver`](@ref) (only for CPU with `Float64`)
-# - [`CGPressureSolver`](@ref)
-# - [`SpectralPressureSolver`](@ref) (only for periodic boundary conditions and
-# uniform grids)
-
-pressure_solver = DirectPressureSolver(setup)
-
-# We will solve for a time interval of ten seconds.
-t_start, t_end = tlims = (T(0), T(10))
-
-# The initial conditions are defined as plain Julia functions.
-initial_velocity_u(x, y) = zero(x)
-initial_velocity_v(x, y) = zero(x)
-initial_pressure(x, y) = zero(x)
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t_start;
- initial_pressure,
- pressure_solver,
-)
+setup = Setup(x, y; boundary_conditions, Re, ArrayType);
+
+# The initial conditions are provided in function. The value `dim()` determines
+# the velocity component.
+ustart = create_initial_conditions(setup, (dim, x, y) -> zero(x));
# ## Solve problems
#
@@ -120,7 +88,8 @@ V₀, p₀ = create_initial_conditions(
# The [`solve_steady_state`](@ref) function is for computing a state where the right hand side of the
# momentum equation is zero.
-V, p = solve_steady_state(setup, V₀, p₀)
+## u, p = solve_steady_state(setup, u₀, p₀)
+nothing
# For this test case, the same steady state may be obtained by solving an
# unsteady problem for a sufficiently long time.
@@ -130,52 +99,47 @@ V, p = solve_steady_state(setup, V₀, p₀)
# later returned by `solve_unsteady`.
processors = (
- field_plotter(device(setup); nupdate = 50),
- ## energy_history_plotter(device(setup); nupdate = 1),
- ## energy_spectrum_plotter(device(setup); nupdate = 100),
- ## animator(device(setup), "vorticity.mkv"; nupdate = 4),
- vtk_writer(setup; nupdate = 100, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1000),
+ ## rtp = realtimeplotter(; setup, plot = fieldplot, nupdate = 50),
+ ## ehist = realtimeplotter(; setup, plot = energy_history_plot, nupdate = 10),
+ ## espec = realtimeplotter(; setup, plot = energy_spectrum_plot, nupdate = 10),
+ ## anim = animator(; setup, path = "$output/solution.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 100, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1000),
);
# By default, a standard fourth order Runge-Kutta method is used. If we don't
# provide the time step explicitly, an adaptive time step is used.
-V, p, outputs =
- solve_unsteady(setup, V₀, p₀, tlims; Δt = T(0.001), processors, pressure_solver, device);
+tlims = (T(0), T(10))
+state, outputs = solve_unsteady(; setup, ustart, tlims, Δt = T(1e-3), processors);
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`
+# We may visualize or export the computed fields
# Export fields to VTK. The file `output/solution.vti` may be opened for
-# visulization in [ParaView](https://www.paraview.org/). This is particularly
+# visualization in [ParaView](https://www.paraview.org/). This is particularly
# useful for inspecting results from 3D simulations.
-save_vtk(setup, V, p, t_end, "output/solution")
+save_vtk(setup, state.u, state.t, "$output/solution")
# Plot pressure
-plot_pressure(setup, p)
+fieldplot(state; setup, fieldname = :pressure)
# Plot velocity. Note the time stamp used for computing boundary conditions, if
# any.
-plot_velocity(setup, V, t_end)
-
-# Plot vorticity (with custom levels)
-levels = [-7, -5, -4, -3, -2, -1, -0.5, 0, 0.5, 1, 2, 3, 7]
-plot_vorticity(setup, V, t_end; levels)
+fieldplot(state; setup, fieldname = :velocity)
-# Plot streamfunction. Note the time stamp used for computing boundary
-# conditions, if any
-plot_streamfunction(setup, V, t_end)
-
-# In addition, the tuple `outputs` contains quantities from our processors.
-#
-# The [`field_plotter`](@ref) returns the field plot figure.
-outputs[1]
-
-# The [`vtk_writer`](@ref) returns the file name of the ParaView collection
-# file. This allows for visualizing the solution time series in ParaView.
-outputs[2]
+# Plot vorticity
+fieldplot(state; setup, fieldname = :vorticity)
+# In addition, the named tuple `outputs` contains quantities from our
+# processors.
# The logger returns nothing.
-outputs[3]
+
+## outputs.rtp
+## outputs.ehist
+## outputs.espec
+## outputs.anim
+## outputs.vtk
+## outputs.field
+outputs.log
diff --git a/examples/LidDrivenCavity3D.jl b/examples/LidDrivenCavity3D.jl
index 356e76cfc..981118614 100644
--- a/examples/LidDrivenCavity3D.jl
+++ b/examples/LidDrivenCavity3D.jl
@@ -1,115 +1,79 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Lid-Driven Cavity - 3D
#
# In this example we consider a box with a moving lid. The velocity is initially at rest. The
# solution should reach at steady state equilibrium after a certain time. The same steady
# state should be obtained when solving a steady state problem.
-# We start by loading packages.
-# A [Makie](https://github.com/JuliaPlots/Makie.jl) plotting backend is needed
-# for plotting. `GLMakie` creates an interactive window (useful for real-time
-# plotting), but does not work when building this example on GitHub.
-# `CairoMakie` makes high-quality static vector-graphics plots.
-
#md using CairoMakie
using GLMakie #!md
using IncompressibleNavierStokes
# Case name for saving results
-name = "LidDrivenCavity3D"
+output = "output/LidDrivenCavity3D"
-# Viscosity model
-viscosity_model = LaminarModel(; Re = 1000.0)
+# Floating point type
+T = Float64
-# Boundary conditions: horizontal movement of the top lid
-u_bc(x, y, z, t) = y ≈ 1.0 ? 1.0 : 0.0
-v_bc(x, y, z, t) = 0.0
-w_bc(x, y, z, t) = y ≈ 1.0 ? 0.2 : 0.0
-bc_type = (;
- u = (;
- x = (:dirichlet, :dirichlet),
- y = (:dirichlet, :dirichlet),
- z = (:periodic, :periodic),
- ),
- v = (;
- x = (:dirichlet, :dirichlet),
- y = (:dirichlet, :dirichlet),
- z = (:periodic, :periodic),
- ),
- w = (;
- x = (:dirichlet, :dirichlet),
- y = (:dirichlet, :dirichlet),
- z = (:periodic, :periodic),
- ),
-)
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+# Reynolds number
+Re = T(1_000)
# A 3D grid is a Cartesian product of three vectors. Here we refine the grid
# near the walls.
-x = cosine_grid(0.0, 1.0, 25)
-y = cosine_grid(0.0, 1.0, 25)
-z = LinRange(-0.2, 0.2, 10)
-plot_grid(x, y, z)
+x = cosine_grid(T(0), T(1), 25)
+y = cosine_grid(T(0), T(1), 25)
+z = LinRange(-T(0.2), T(0.2), 11)
+plotgrid(x, y, z)
-# Build setup and assemble operators
-setup = Setup(x, y, z; viscosity_model, u_bc, v_bc, w_bc, bc_type);
+# Boundary conditions: horizontal movement of the top lid
+U(dim, x, y, z, t) = dim() == 1 ? one(x) : dim() == 2 ? zero(x) : one(x) / 5
+dUdt(dim, x, y, z, t) = zero(x)
+boundary_conditions = (
+ ## x left, x right
+ (DirichletBC(), DirichletBC()),
-# Time interval
-t_start, t_end = tlims = (0.0, 0.2)
+ ## y rear, y front
+ (DirichletBC(), DirichletBC(U, dUdt)),
-# Initial conditions
-initial_velocity_u(x, y, z) = 0.0
-initial_velocity_v(x, y, z) = 0.0
-initial_velocity_w(x, y, z) = 0.0
-initial_pressure(x, y, z) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- initial_velocity_w,
- t_start;
- initial_pressure,
-);
+ ## z bottom, z top
+ (PeriodicBC(), PeriodicBC()),
+)
-# Solve steady state problem
-V, p = solve_steady_state(setup, V₀, p₀; npicard = 5, maxiter = 15);
+# Build setup and assemble operators
+setup = Setup(x, y, z; Re, boundary_conditions, ArrayType);
-# Iteration processors
-processors = (
- field_plotter(setup; nupdate = 1),
- ## energy_history_plotter(setup; nupdate = 1),
- ## energy_spectrum_plotter(setup; nupdate = 100),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 5, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
-);
+# Initial conditions
+ustart = create_initial_conditions(setup, (dim, x, y, z) -> zero(x))
# Solve unsteady problem
-V, p, outputs = solve_unsteady(setup, V₀, p₀, tlims; Δt = 0.001, processors)
-#md current_figure()
+(; u, t), outputs = solve_unsteady(;
+ setup,
+ ustart,
+ tlims = (T(0), T(0.2)),
+ Δt = T(1e-3),
+ processors = (
+ ## rtp = realtimeplotter(; setup, plot = fieldplot, nupdate = 50),
+ ehist = realtimeplotter(; setup, plot = energy_history_plot, nupdate = 10),
+ ## espec = realtimeplotter(; setup, plot = energy_spectrum_plot, nupdate = 10),
+ ## anim = animator(; setup, path = "$output/solution.mkv", nupdate = 20),
+ # vtk = vtk_writer(; setup, nupdate = 100, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 20),
+ ),
+);
# ## Post-process
#
# We may visualize or export the computed fields `(V, p)`
# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
-
-# Plot pressure
-plot_pressure(setup, p)
-
-# Plot velocity
-plot_velocity(setup, V, t_end)
-
-# Plot vorticity
-plot_vorticity(setup, V, t_end)
+save_vtk(setup, u, t, "$output/solution")
-# Plot streamfunction
-## plot_streamfunction(setup, V, t_end)
-nothing
+# Energy history
+outputs.ehist
diff --git a/examples/MultiActuator.jl b/examples/MultiActuator.jl
new file mode 100644
index 000000000..99ea04cf7
--- /dev/null
+++ b/examples/MultiActuator.jl
@@ -0,0 +1,158 @@
+# # Unsteady actuator case - 2D
+#
+# In this example, an unsteady inlet velocity profile at encounters a wind
+# turbine blade in a wall-less domain. The blade is modeled as a uniform body
+# force on a thin rectangle.
+
+#md using CairoMakie
+using GLMakie #!md
+using IncompressibleNavierStokes
+using Random
+
+# Output directory
+output = "output/MultiActuator"
+
+# Floating point precision
+T = Float64
+
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+using CUDA;
+T = Float32;
+# T = Float64;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+
+set_theme!(; GLMakie = (; scalefactor = 1.5))
+
+const ππ = T(π)
+
+# Boundary conditions
+boundary_conditions = (
+ ## x left, x right
+ (
+ ## Unsteady BC requires time derivatives
+ DirichletBC(
+ (dim, x, y, t) -> sin(ππ / 6 * sin(ππ / 6 * t) + ππ / 2 * (dim() == 1)),
+ (dim, x, y, t) ->
+ (ππ / 6)^2 *
+ cos(ππ / 6 * t) *
+ cos(ππ / 6 * sin(ππ / 6 * t) + ππ / 2 * (dim() == 1)),
+ ),
+ PressureBC(),
+ ),
+
+ ## y rear, y front
+ (PressureBC(), PressureBC()),
+)
+
+# Actuator body force: A thrust coefficient `Cₜ` distributed over a thin rectangle
+create_bodyforce(; xc, yc, D, δ, C) =
+ (dim, x, y, t) ->
+ dim() == 1 && abs(x - xc) ≤ δ / 2 && abs(y - yc) ≤ D / 2 ? -C / (D * δ) : zero(x)
+
+create_manyforce(forces...) = function (dim, x, y, t)
+ out = zero(x)
+ for f in forces
+ out += f(dim, x, y, t)
+ end
+ out
+end
+
+disk = (; D = T(1), δ = T(0.11), C = T(0.2))
+bodyforce = create_manyforce(
+ create_bodyforce(; xc = T(2), yc = T(0), disk...),
+ create_bodyforce(; xc = T(4), yc = T(0.7), disk...),
+ create_bodyforce(; xc = T(6.4), yc = T(-1), disk...),
+)
+
+# A 2D grid is a Cartesian product of two vectors
+n = 80
+x = LinRange(T(0), T(10), 5n + 1)
+y = LinRange(-T(2), T(2), 2n + 1)
+
+# Build setup and assemble operators
+setup = Setup(x, y; Re = T(2000), boundary_conditions, bodyforce, ArrayType);
+
+# Initial conditions (extend inflow)
+ustart = create_initial_conditions(setup, (dim, x, y) -> dim() == 1 ? one(x) : zero(x));
+t = T(0)
+
+# # We create a box to visualize the actuator.
+# (; xc, yc, D, δ) = setup.bodyforce
+# box = [
+# Point2f(xc - δ / 2, yc + D / 2),
+# Point2f(xc - δ / 2, yc - D / 2),
+# Point2f(xc + δ / 2, yc - D / 2),
+# Point2f(xc + δ / 2, yc + D / 2),
+# Point2f(xc - δ / 2, yc + D / 2),
+# ]
+
+boxes = map(bodyforce.forces) do (; xc, yc, D, δ)
+ [
+ Point2f(xc - δ / 2, yc + D / 2),
+ Point2f(xc - δ / 2, yc - D / 2),
+ Point2f(xc + δ / 2, yc - D / 2),
+ Point2f(xc + δ / 2, yc + D / 2),
+ Point2f(xc - δ / 2, yc + D / 2),
+ ]
+end
+box = boxes[1]
+
+# Solve unsteady problem
+state, outputs = solve_unsteady(;
+ setup,
+ ustart,
+ tlims = (T(0), 4 * T(12)),
+ # (T(0), T(1));
+ method = RKMethods.RK44P2(),
+ Δt = T(0.01),
+ processors = (
+ rtp = realtimeplotter(;
+ setup,
+ # plot = fieldplot,
+ # fieldname = :velocity,
+ # fieldname = :pressure,
+ nupdate = 1,
+ ),
+ boxplotter = processor() do state
+ for box in boxes
+ lines!(current_axis(), box; color = :red)
+ end
+ end,
+ ## ehist = realtimeplotter(; setup, plot = energy_history_plot, nupdate = 1),
+ ## espec = realtimeplotter(; setup, plot = energy_spectrum_plot, nupdate = 1),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = "$output", filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
+);
+
+# ## Post-process
+#
+# We may visualize or export the computed fields `(u, p)`.
+
+# Export to VTK
+save_vtk(setup, state.u, state.t, "$output/solution")
+
+# Plot pressure
+fig = fieldplot(state; setup, fieldname = :pressure)
+# lines!(box...; color = :red)
+lines!.(boxes; color = :red);
+fig
+
+# Plot velocity
+fig = fieldplot(state; setup, fieldname = :velocity)
+lines!.(boxes; color = :red);
+fig
+
+# Plot vorticity
+fig = fieldplot(state; setup, fieldname = :vorticity)
+lines!.(boxes; color = :red);
+fig
diff --git a/examples/PlanarMixing2D.jl b/examples/PlanarMixing2D.jl
index 1c1508eb4..97d3ef8f7 100644
--- a/examples/PlanarMixing2D.jl
+++ b/examples/PlanarMixing2D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Planar mixing - 2D
#
# Planar mixing example, as presented in [List2022](@cite).
@@ -19,29 +12,30 @@ end #src
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "PlanarMixing2D"
+# Output directory
+output = "output/PlanarMixing2D"
# Viscosity model
-viscosity_model = LaminarModel(; Re = 500.0)
+Re = 500.0
# Boundary conditions: Unsteady BC requires time derivatives
ΔU = 1.0
-Ū = 1.0
-ϵ = (0.082Ū, 0.012Ū)
+Ubar = 1.0
+ϵ = (0.082Ubar, 0.012Ubar)
n = (0.4π, 0.3π)
ω = (0.22, 0.11)
-u_bc(x, y, t) =
- x ≈ 0.0 ?
+U(dim, x, y, t) =
+ dim() == 1 ?
1.0 + ΔU / 2 * tanh(2y) + sum(@. ϵ * (1 - tanh(y / 2)^2) * cos(n * y) * sin(ω * t)) :
0.0
-v_bc(x, y, t) = 0.0
-dudt_bc(x, y, t) =
- x ≈ 0.0 ? sum(@. ϵ * (1 - tanh(y / 2)^2) * cos(n * y) * ω * cos(ω * t)) : 0.0
-dvdt_bc(x, y, t) = 0.0
-bc_type = (;
- u = (; x = (:dirichlet, :pressure), y = (:symmetric, :symmetric)),
- v = (; x = (:dirichlet, :symmetric), y = (:pressure, :pressure)),
+dUdt(dim, x, y, t) =
+ dim() == 1 ? sum(@. ϵ * (1 - tanh(y / 2)^2) * cos(n * y) * ω * cos(ω * t)) : 0.0
+boundary_conditions = (
+ ## x left, x right
+ (DirichletBC(U, dUdt), PressureBC()),
+
+ ## y rear, y front
+ (PressureBC(), PressureBC()),
)
# A 2D grid is a Cartesian product of two vectors
@@ -49,65 +43,43 @@ n = 64
## n = 256
x = LinRange(0.0, 256.0, 4n)
y = LinRange(-32.0, 32.0, n)
-plot_grid(x, y)
+plotgrid(x, y)
# Build setup and assemble operators
-setup = Setup(x, y; viscosity_model, u_bc, v_bc, dudt_bc, dvdt_bc, bc_type);
-
-# Time interval
-t_start, t_end = tlims = (0.0, 100.0)
-
-# Initial conditions
-initial_velocity_u(x, y) = u_bc(0.0, y, 0.0)
-initial_velocity_v(x, y) = 0.0
-initial_pressure(x, y) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t_start;
- initial_pressure,
-);
+setup = Setup(x, y; Re, boundary_conditions);
+psolver = psolver_direct(setup);
-# Iteration processors
-processors = (
- field_plotter(setup; nupdate = 1),
- ## energy_history_plotter(setup; nupdate = 1),
- ## energy_spectrum_plotter(setup; nupdate = 100),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
-);
+# Initial conditions (extend inflow)
+ustart = create_initial_conditions(setup, (dim, x, y) -> U(dim, x, y, 0.0); psolver);
# Solve unsteady problem
-V, p, outputs = solve_unsteady(
+state, outputs = solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- method = RK44P2(),
+ ustart,
+ tlims = (0.0, 100.0),
+ psolver,
+ method = RKMethods.RK44P2(),
Δt = 0.1,
- processors,
- inplace = true,
+ processors = (
+ rtp = realtimeplotter(;
+ setup,
+ plot = fieldplot,
+ ## plot = energy_history_plot,
+ ## plot = energy_spectrum_plot,
+ nupdate = 1,
+ ),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
);
-#md current_figure()
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`
+# We may visualize or export the computed fields `(u, p)`
-# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
-
-# Plot pressure
-plot_pressure(setup, p)
-
-# Plot velocity
-plot_velocity(setup, V, t_end)
+outputs.rtp
-# Plot vorticity
-plot_vorticity(setup, V, t_end)
-
-# Plot streamfunction
-plot_streamfunction(setup, V, t_end)
+# Export to VTK
+save_vtk(setup, state.u, state.t, "$output/solution"; psolver)
diff --git a/examples/PlaneJets2D.jl b/examples/PlaneJets2D.jl
index 87373b041..2d01473c8 100644
--- a/examples/PlaneJets2D.jl
+++ b/examples/PlaneJets2D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Plane jets - 2D
#
# Plane jets example, as presented in [MacArt2021](@cite). Note that the
@@ -22,184 +15,175 @@ using GLMakie #!md
using IncompressibleNavierStokes
using LaTeXStrings
-# Case name for saving results
-name = "PlaneJets2D"
-
-# Viscosity model
-viscosity_model = LaminarModel(; Re = 6000.0)
-
-# Test cases (A, B, C, D; in order)
-# U() = sqrt(467.4)
-U() = 21.619435700313733
+# Output directory
+output = "output/PlaneJets2D"
-u_A(y) = U() / 2 * (tanh((y + 1 / 2) / 0.1) - tanh((y - 1 / 2) / 0.1))
+# Floating point type
+T = Float64
-u_B(y) =
- U() / 2 * (tanh((y + 1 + 1 / 2) / 0.1) - tanh((y + 1 - 1 / 2) / 0.1)) +
- U() / 2 * (tanh((y - 1 + 1 / 2) / 0.1) - tanh((y - 1 - 1 / 2) / 0.1))
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
-u_C(y) =
- U() / 2 * (tanh(((y + 1.0) / 1 + 1 / 2) / 0.1) - tanh(((y + 1.0) / 1 - 1 / 2) / 0.1)) +
- U() / 4 * (tanh(((y - 1.5) / 2 + 1 / 2) / 0.2) - tanh(((y - 1.5) / 2 - 1 / 2) / 0.2))
+# Reynolds number
+Re = T(6_000)
-u_D(y) =
- U() / 2 * (tanh(((y + 1.0) / 1 + 1 / 2) / 0.1) - tanh(((y + 1.0) / 1 - 1 / 2) / 0.1)) -
- U() / 4 * (tanh(((y - 1.5) / 2 + 1 / 2) / 0.2) - tanh(((y - 1.5) / 2 - 1 / 2) / 0.2))
-
-# u(y) = u_A(y)
-# u(y) = u_B(y)
-u(y) = u_C(y)
-# u(y) = u_D(y)
+# Test cases (A, B, C, D; in order)
+## V() = sqrt(T(467.4))
+V() = T(21.619435700313733)
+
+U_A(y) = V() / 2 * (tanh((y + T(0.5)) / T(0.1)) - tanh((y - T(0.5)) / T(0.1)))
+
+U_B(y) =
+ V() / 2 * (tanh((y + 1 + T(0.5)) / T(0.1)) - tanh((y + 1 - T(0.5)) / T(0.1))) +
+ V() / 2 * (tanh((y - 1 + T(0.5)) / T(0.1)) - tanh((y - 1 - T(0.5)) / T(0.1)))
+
+U_C(y) =
+ V() / 2 * (
+ tanh(((y + T(1.0)) / 1 + T(0.5)) / T(0.1)) -
+ tanh(((y + T(1.0)) / 1 - T(0.5)) / T(0.1))
+ ) +
+ V() / 4 * (
+ tanh(((y - T(1.5)) / 2 + T(0.5)) / T(0.2)) -
+ tanh(((y - T(1.5)) / 2 - T(0.5)) / T(0.2))
+ )
+
+U_D(y) =
+ V() / 2 * (
+ tanh(((y + T(1.0)) / 1 + T(0.5)) / T(0.1)) -
+ tanh(((y + T(1.0)) / 1 - T(0.5)) / T(0.1))
+ ) -
+ V() / 4 * (
+ tanh(((y - T(1.5)) / 2 + T(0.5)) / T(0.2)) -
+ tanh(((y - T(1.5)) / 2 - T(0.5)) / T(0.2))
+ )
+
+## U(y) = U_A(y)
+## U(y) = U_B(y)
+U(y) = U_C(y)
+## U(y) = U_D(y)
# Random noise to stimulate turbulence
-u(x, y) = (1 + 0.1 * (rand() - 1 / 2)) * u(y)
+U(x, y) = (1 + T(0.1) * (rand(T) - T(0.5))) * U(y)
-# # Boundary conditions: Unsteady BC requires time derivatives
-# u_bc(x, y, t) = x ≈ 0.0 ? u(x, y) : 0.0
-# v_bc(x, y, t) = 0.0
-# bc_type = (;
-# u = (; x = (:periodic, :periodic), y = (:symmetric, :symmetric)),
-# v = (; x = (:periodic, :periodic), y = (:pressure, :pressure)),
-# )
+## boundary_conditions = (
+## (PeriodicBC(), PeriodicBC()),
+## (PressureBC(), PressureBC())
+## )
# A 2D grid is a Cartesian product of two vectors
n = 64
## n = 128
## n = 256
-x = LinRange(0.0, 16.0, 4n)
-y = LinRange(-10.0, 10.0, 5n)
-plot_grid(x, y)
+x = LinRange(T(0), T(16), 4n + 1)
+y = LinRange(-T(10), T(10), 5n + 1)
+plotgrid(x, y)
# Build setup and assemble operators
-setup = Setup(x, y; viscosity_model);
-# setup = Setup(x, y; viscosity_model, u_bc, v_bc, bc_type);
-
-# Since the grid is uniform and identical for x and y, we may use a specialized
-# spectral pressure solver
-pressure_solver = SpectralPressureSolver(setup)
-
-# Time interval
-t_start, t_end = tlims = (0.0, 1.0)
+setup = Setup(x, y; Re, ArrayType);
+## setup = Setup(x, y; Re, boundary_conditions, ArrayType);
# Initial conditions
-initial_velocity_u(x, y) = u(x, y)
-initial_velocity_v(x, y) = 0.0
-initial_pressure(x, y) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t_start;
- initial_pressure,
- pressure_solver,
-);
-V, p = V₀, p₀
+ustart = create_initial_conditions(setup, (dim, x, y) -> dim() == 1 ? U(x, y) : zero(x));
# Real time plot: Streamwise average and spectrum
-mean_plotter(setup; nupdate = 1) = processor(
- function (state)
- (; indu, yu, yin, Nux_in, Nuy_in) = setup.grid
-
- umean = @lift begin
- (; V, p, t) = $state
- u = V[indu]
- sleep(0.001)
- reshape(sum(reshape(u, size(yu)); dims = 1), :) ./ (Nux_in * U())
- end
-
- K = Nux_in ÷ 2
- k = 1:(K-1)
-
- # Find energy spectrum where y = 0
- n₀ = Nuy_in ÷ 2
- E₀ = @lift begin
- (; V, p, t) = $state
- u = V[indu]
- u_y = reshape(u, size(yu))[:, n₀]
- abs.(fft(u_y .^ 2))[k.+1]
- end
-
- # Find energy spectrum where y = 1
- n₁ = argmin(n -> abs(yin[n] .- 1), 1:Nuy_in)
- E₁ = @lift begin
- (; V, p, t) = $state
- u = V[indu]
- u_y = reshape(u, size(yu))[:, n₁]
- abs.(fft(u_y .^ 2))[k.+1]
- end
-
- fig = Figure()
- ax = Axis(
- fig[1, 1];
- title = "Mean streamwise flow",
- xlabel = "y",
- ylabel = L"\langle u \rangle / U_0",
- )
- lines!(ax, yu[1, :], umean)
- ax = Axis(
- fig[1, 2];
- title = "Streamwise energy spectrum",
- xscale = log10,
- yscale = log10,
- xlabel = L"k_x",
- ylabel = L"\hat{U}_{cl} / U_0",
- )
- # ylims!(ax, (10^(0.0), 10^4.0))
- ksub = k[10:end]
- lines!(ax, ksub, 1000 .* ksub .^ (-3 / 5); label = L"k^{-3/5}")
- lines!(ax, ksub, 1e7 .* ksub .^ -3; label = L"k^{-3}")
- scatter!(ax, k, E₀; label = "y = $(yin[n₀])")
- scatter!(ax, k, E₁; label = "y = $(yin[n₁])")
- axislegend(ax; position = :lb)
-
- display(fig)
- fig
- end;
- nupdate,
-)
-
-# Iteration processors
-processors = (
- field_plotter(setup; nupdate = 1),
- ## energy_history_plotter(setup; nupdate = 1),
- ## energy_spectrum_plotter(setup; nupdate = 100),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
- mean_plotter(setup),
-);
+function meanplot(state; setup)
+ (; xp, Iu, Ip, Nu, N) = setup.grid
+
+ umean = lift(state) do (; u, p, t)
+ reshape(sum(u[1][Iu[1]]; dims = 1), :) ./ Nu[1][1] ./ V()
+ end
+
+ K = Nu[1][2] ÷ 2
+ k = 1:(K-1)
+
+ ## Find energy spectrum where y = 0
+ n₀ = findmin(abs, xp[2])[2]
+ E₀ = lift(state) do (; u, p, t)
+ u_y = u[1][:, n₀]
+ abs.(fft(u_y .^ 2))[k.+1]
+ end
+ y₀ = xp[2][n₀]
+
+ ## Find energy spectrum where y = 1
+ n₁ = findmin(y -> abs(y - 1), xp[2])[2]
+ E₁ = lift(state) do (; u, p, t)
+ u_y = u[1][:, n₁]
+ abs.(fft(u_y .^ 2))[k.+1]
+ end
+ y₁ = xp[2][n₁]
+
+ fig = Figure()
+ ax = Axis(
+ fig[1, 1];
+ title = "Mean streamwise flow",
+ xlabel = "y",
+ ylabel = L"\langle u \rangle / U_0",
+ )
+ lines!(ax, xp[2][2:end-1], umean)
+ ax = Axis(
+ fig[1, 2];
+ title = "Streamwise energy spectrum",
+ xscale = log10,
+ yscale = log10,
+ xlabel = L"k_x",
+ ylabel = L"\hat{U}_{cl} / U_0",
+ )
+ ## ylims!(ax, (10^(0.0), 10^4.0))
+ ksub = k[10:end]
+ ## lines!(ax, ksub, 1000 .* ksub .^ (-5 / 3); label = L"k^{-5/3}")
+ lines!(ax, ksub, 1e7 .* ksub .^ -3; label = L"k^{-3}")
+ scatter!(ax, k, E₀; label = "y = $y₀")
+ scatter!(ax, k, E₁; label = "y = $y₁")
+ axislegend(ax; position = :lb)
+ ## on(_ -> autolimits!(ax), E₁)
+
+ fig
+end
# Solve unsteady problem
-V, p, outputs = solve_unsteady(
+state, outputs = solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- method = RK44P2(),
+ ustart,
+ tlims = (T(0), T(1)),
+ method = RKMethods.RK44P2(),
Δt = 0.001,
- processors,
- pressure_solver,
- inplace = true,
+ processors = (
+ rtp = realtimeplotter(;
+ setup,
+ ## plot = fieldplot,
+ ## plot = energy_history_plot,
+ ## plot = energy_spectrum_plot,
+ plot = meanplot,
+ nupdate = 1,
+ ),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 4),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
);
-#md current_figure()
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`
+# We may visualize or export the computed fields `(u, p)`
+
+outputs.rtp
# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
+save_vtk(setup, state.u, state.t, "$output/solution")
# Plot pressure
-plot_pressure(setup, p)
+fieldplot(state; setup, fieldname = :pressure)
-# Plot velocity
-plot_velocity(setup, V₀, t_end)
-plot_velocity(setup, V, t_end)
+# Plot initial velocity
+fieldplot((; u = u₀, p = p₀, t = T(0)); setup, fieldname = :velocity)
-# Plot vorticity
-plot_vorticity(setup, V, t_end)
+# Plot final velocity
+fieldplot(state; setup, fieldname = :velocity)
-# Plot stream function
-plot_streamfunction(setup, V, t_end)
+# Plot vorticity
+fieldplot(state; setup, fieldname = :vorticity)
diff --git a/examples/Project.toml b/examples/Project.toml
index 2fe454328..6afb14999 100644
--- a/examples/Project.toml
+++ b/examples/Project.toml
@@ -1,6 +1,17 @@
+name = "Examples"
+uuid = "318dbb63-4243-420f-99f2-d56058123f9d"
+authors = ["Syver Døving Agdestein "]
+version = "0.1.0"
+
[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e"
CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
-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"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
diff --git a/examples/README.md b/examples/README.md
new file mode 100644
index 000000000..82a545fb9
--- /dev/null
+++ b/examples/README.md
@@ -0,0 +1,25 @@
+# Examples
+
+Example scripts for
+[IncompressibleNavierStokes.jl](https://github.com/agdestein/IncompressibleNavierStokes.jl).
+
+## Set up environment
+
+From this directory, run:
+
+```sh
+julia --project -e '
+using Pkg
+Pkg.develop(PackageSpec(; path = ".."))
+Pkg.instantiate()
+'
+```
+
+or interactively from a Julia REPL:
+
+```julia-repl
+julia> ]
+(v1.10) pkg> activate .
+(Examples) pkg> dev ..
+(Examples) pkg> instantiate
+```
diff --git a/examples/RayleighBenard2D.jl b/examples/RayleighBenard2D.jl
new file mode 100644
index 000000000..5bcb6452f
--- /dev/null
+++ b/examples/RayleighBenard2D.jl
@@ -0,0 +1,57 @@
+#md using CairoMakie
+using GLMakie #!md
+using IncompressibleNavierStokes
+
+using CUDA, CUDSS
+ArrayType = CuArray
+
+outdir = joinpath(@__DIR__, "output")
+
+temperature = temperature_equation(;
+ Pr = T(0.71),
+ Ra = T(1e7),
+ Ge = T(0.1),
+ dodissipation = true,
+ boundary_conditions = (
+ (SymmetricBC(), SymmetricBC()),
+ (DirichletBC(1.0), DirichletBC(0.0)),
+ ),
+ gdir = 2,
+ nondim_type = 1,
+)
+
+n = 64
+x = LinRange(T(0), T(1), n + 1)
+y = LinRange(T(0), T(1), n + 1)
+setup = Setup(
+ x,
+ y;
+ boundary_conditions = ((DirichletBC(), DirichletBC()), (DirichletBC(), DirichletBC())),
+ Re = 1 / temperature.α1,
+ temperature,
+ ArrayType,
+);
+ustart = create_initial_conditions(setup, (dim, x, y) -> zero(x));
+(; xp) = setup.grid;
+tempstart = @. max(-sin($(T(3π)) * xp[1]) / 100, 0) + (1 - xp[2]');
+
+state, outputs = solve_unsteady(;
+ setup,
+ ustart,
+ tempstart,
+ tlims = (0.0, 100.0),
+ Δt = 1e-2,
+ processors = (;
+ # rtp = realtimeplotter(;
+ anim = animator(;
+ path = "$outdir/RB2D.mp4",
+ setup,
+ nupdate = 100,
+ fieldname = :temperature,
+ colorrange = (0.0, 1.0),
+ # displayupdates = true,
+ size = (600, 500),
+ ),
+ log = timelogger(; nupdate = 20),
+ ),
+);
diff --git a/examples/RayleighTaylor2D.jl b/examples/RayleighTaylor2D.jl
new file mode 100644
index 000000000..5d294ed61
--- /dev/null
+++ b/examples/RayleighTaylor2D.jl
@@ -0,0 +1,69 @@
+#md using CairoMakie
+using GLMakie #!md
+using CairoMakie
+using IncompressibleNavierStokes
+
+using CUDA, CUDSS
+T = Float32
+ArrayType = CuArray
+
+outdir = joinpath(@__DIR__, "output")
+
+temperature = temperature_equation(;
+ Pr = T(0.71),
+ Ra = T(1e6),
+ Ge = T(1.0),
+ # Ge = T(0.1),
+ dodissipation = true,
+ boundary_conditions = ((SymmetricBC(), SymmetricBC()), (SymmetricBC(), SymmetricBC())),
+ gdir = 2,
+ nondim_type = 1,
+)
+
+n = 128
+x = LinRange(T(0), T(1), n + 1)
+y = LinRange(T(0), T(2), 2n + 1)
+setup = Setup(
+ x,
+ y;
+ boundary_conditions = ((DirichletBC(), DirichletBC()), (DirichletBC(), DirichletBC())),
+ Re = 1 / temperature.α1,
+ temperature,
+ ArrayType,
+);
+ustart = create_initial_conditions(setup, (dim, x, y) -> zero(x));
+(; xp) = setup.grid;
+tempstart = @. $(T(1)) * (1.0 + 0.05 * sin($(T(2π)) * xp[1]) > xp[2]');
+tempstart = @. $(T(1)) * (1.0 + 0.05 * sin($(T(π)) * xp[1]) > xp[2]');
+# @. tempstart += 0.3 * randn()
+
+state, outputs = solve_unsteady(;
+ setup,
+ # ustart,
+ # tempstart,
+ ustart = state.u,
+ tempstart = state.temp,
+ tlims = (0.0, 40.0),
+ Δt = 5e-3,
+ processors = (;
+ rtp = realtimeplotter(;
+ # anim = animator(;
+ path = "$outdir/RT2D.mp4",
+ setup,
+ nupdate = 200,
+ fieldname = :temperature,
+ displayupdates = true,
+ size = (400, 600),
+ ),
+ log = timelogger(; nupdate = 50),
+ ),
+);
+
+state.u .|> extrema
+state.temp |> extrema
+state.u[1]
+state.u[2]
+state.temp
+
+fieldplot(state; setup, fieldname = :temperature)
+save("toto.png", current_figure())
diff --git a/examples/RayleighTaylor3D.jl b/examples/RayleighTaylor3D.jl
new file mode 100644
index 000000000..a425c3f1b
--- /dev/null
+++ b/examples/RayleighTaylor3D.jl
@@ -0,0 +1,96 @@
+#md using CairoMakie
+using GLMakie #!md
+using IncompressibleNavierStokes
+
+using CUDA, CUDSS
+T = Float32
+ArrayType = CuArray
+
+outdir = joinpath(@__DIR__, "output")
+
+temperature = temperature_equation(;
+ Pr = T(0.71),
+ Ra = T(1e6),
+ # Ge = T(0.1),
+ Ge = T(1.0),
+ dodissipation = true,
+ boundary_conditions = (
+ (SymmetricBC(), SymmetricBC(), SymmetricBC()),
+ (SymmetricBC(), SymmetricBC(), SymmetricBC()),
+ (SymmetricBC(), SymmetricBC(), SymmetricBC()),
+ ),
+ gdir = 3,
+ nondim_type = 1,
+)
+
+n = 80
+x = LinRange(T(0), T(1), n + 1)
+y = LinRange(T(0), T(1), n + 1)
+z = LinRange(T(0), T(2), 2n + 1)
+setup = Setup(
+ x,
+ y,
+ z;
+ boundary_conditions = (
+ (DirichletBC(), DirichletBC(), DirichletBC()),
+ (DirichletBC(), DirichletBC(), DirichletBC()),
+ (DirichletBC(), DirichletBC(), DirichletBC()),
+ ),
+ Re = 1 / temperature.α1,
+ temperature,
+ ArrayType,
+);
+
+psolver = psolver_direct(setup)
+
+ustart = create_initial_conditions(setup, (dim, x, y, z) -> zero(x); psolver);
+(; xp) = setup.grid;
+xx = xp[1];
+xy = reshape(xp[2], 1, :);
+xz = reshape(xp[3], 1, 1, :);
+tempstart = @. $(T(1)) * (1.0 + 0.05 * sin($(T(1.05 * π)) * xx) * sin($(T(π)) * xy) > xz);
+
+fieldplot(
+ (; u = ustart, temp = tempstart, t = T(0));
+ # state;
+ setup,
+ levels = LinRange{T}(0.8, 1, 5),
+ # levels = LinRange(-T(5), T(1), 10),
+ # fieldname = :eig2field,
+ fieldname = :temperature,
+ size = (400, 600),
+)
+
+save("$outdir/RT3D_initial.png", current_figure())
+
+state, outputs = solve_unsteady(;
+ setup,
+ ustart,
+ tempstart,
+ tlims = (0.0, 40.0),
+ # Δt = 5e-3,
+ Δt = 1e-2,
+ psolver,
+ processors = (;
+ # rtp = realtimeplotter(;
+ anim = animator(;
+ # path = "$outdir/RT3D_temp.mp4",
+ path = "$outdir/RT3D_l2.mp4",
+ setup,
+ nupdate = 20,
+ fieldname = :eig2field,
+ # fieldname = :temperature,
+ levels = LinRange(-T(5), T(1), 10),
+ # levels = LinRange{T}(0, 1, 10),
+ # displayupdates = true,
+ size = (400, 600),
+ ),
+ log = timelogger(; nupdate = 10),
+ ),
+);
+
+field = IncompressibleNavierStokes.eig2field(state.u, setup)[setup.grid.Ip]
+# hist(vec(Array(log(max(eps(T), field)))
+hist(vec(Array(log.(max.(eps(T), .-field)))))
+
+fieldplot(state; setup, fieldname = :temperature)
diff --git a/examples/ShearLayer2D.jl b/examples/ShearLayer2D.jl
index 8215abe5d..e0870590e 100644
--- a/examples/ShearLayer2D.jl
+++ b/examples/ShearLayer2D.jl
@@ -1,10 +1,3 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Shear layer - 2D
#
# Shear layer example.
@@ -19,80 +12,75 @@ end #src
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "ShearLayer2D"
+# Output directory
+output = "output/ShearLayer2D"
+
+# Floating point type
+T = Float64
+
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
-# Viscosity model
-viscosity_model = LaminarModel(; Re = Inf)
+# Reynolds number
+Re = T(Inf)
# A 2D grid is a Cartesian product of two vectors
-n = 100
-x = LinRange(0, 2π, n + 1)
-y = LinRange(0, 2π, n + 1)
-plot_grid(x, y)
+n = 128
+lims = T(0), T(2π)
+x = LinRange(lims..., n + 1)
+y = LinRange(lims..., n + 1)
+plotgrid(x, y)
# Build setup and assemble operators
-setup = Setup(x, y; viscosity_model);
-
-# Time interval
-t_start, t_end = tlims = (0.0, 8.0)
+setup = Setup(x, y; Re, ArrayType);
# Initial conditions: We add 1 to u in order to make global momentum
# conservation less trivial
-d = π / 15
-e = 0.05
-initial_velocity_u(x, y) = y ≤ π ? tanh((y - π / 2) / d) : tanh((3π / 2 - y) / d)
-## initial_velocity_u(x, y) = 1.0 + (y ≤ π ? tanh((y - π / 2) / d) : tanh((3π / 2 - y) / d))
-initial_velocity_v(x, y) = e * sin(x)
-initial_pressure(x, y) = 0.0
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t_start;
- initial_pressure,
-);
-
-# Iteration processors
-processors = (
- field_plotter(setup; nupdate = 1),
- ## energy_history_plotter(setup; nupdate = 1),
- ## energy_spectrum_plotter(setup; nupdate = 100),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
-);
+d = T(π / 15)
+e = T(0.05)
+U1(y) = y ≤ π ? tanh((y - T(π / 2)) / d) : tanh((T(3π / 2) - y) / d)
+## U1(y) = T(1) + (y ≤ π ? tanh((y - T(π / 2)) / d) : tanh((T(3π / 2) - y) / d))
+ustart = create_initial_conditions(setup, (dim, x, y) -> dim() == 1 ? U1(y) : e * sin(x));
# Solve unsteady problem
-V, p, outputs = solve_unsteady(
+state, outputs = solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- method = RK44(),
- Δt = 0.01,
- processors,
- inplace = true,
+ ustart,
+ tlims = (T(0), T(8)),
+ Δt = T(0.01),
+ processors = (
+ rtp = realtimeplotter(;
+ setup,
+ plot = fieldplot,
+ ## plot = energy_history_plot,
+ ## plot = energy_spectrum_plot,
+ nupdate = 1,
+ ),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
);
-#md current_figure()
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`
+# We may visualize or export the computed fields
+
+outputs.rtp
# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
+save_vtk(setup, state.u, state.t, "$output/solution")
# Plot pressure
-plot_pressure(setup, p)
+fieldplot(state; setup, fieldname = :pressure)
# Plot velocity
-plot_velocity(setup, V, t_end)
+fieldplot(state; setup, fieldname = :velocity)
# Plot vorticity
-plot_vorticity(setup, V, t_end)
-
-# Plot streamfunction
-## plot_streamfunction(setup, V, t_end)
-nothing
+fieldplot(state; setup, fieldname = :vorticity)
diff --git a/examples/TaylorGreenVortex2D.jl b/examples/TaylorGreenVortex2D.jl
index 4b8e7399b..4e4898771 100644
--- a/examples/TaylorGreenVortex2D.jl
+++ b/examples/TaylorGreenVortex2D.jl
@@ -1,109 +1,86 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
-# # Taylor-Green vortex - 2D
+# # Convergence study: Taylor-Green vortex (2D)
#
# In this example we consider the Taylor-Green vortex.
-
-# We start by loading packages.
-# A [Makie](https://github.com/JuliaPlots/Makie.jl) plotting backend is needed
-# for plotting. `GLMakie` creates an interactive window (useful for real-time
-# plotting), but does not work when building this example on GitHub.
-# `CairoMakie` makes high-quality static vector-graphics plots.
+# In 2D, it has an analytical solution, given by
+#
+# ```math
+# \begin{split}
+# u^1(x, y, t) & = - \sin(x) \cos(y) \mathrm{e}^{-2 t / Re} \\
+# u^2(x, y, t) & = + \cos(x) \sin(y) \mathrm{e}^{-2 t / Re}
+# \end{split}
+# ```
+#
+# This allows us to test the convergence of our solver.
#md using CairoMakie
using GLMakie #!md
using IncompressibleNavierStokes
-
-# Case name for saving results
-name = "TaylorGreenVortex2D"
-
-# Floating point type
-T = Float32
-
-# Viscosity model
-viscosity_model = LaminarModel(; Re = T(2_000))
-
-# A 2D grid is a Cartesian product of two vectors
-n = 128
-lims = (T(0), T(2π))
-x = LinRange(lims..., n + 1)
-y = LinRange(lims..., n + 1)
-plot_grid(x, y)
-
-# Build setup and assemble operators
-setup = Setup(x, y; viscosity_model);
-
-# Since the grid is uniform and identical for x and y, we may use a specialized
-# spectral pressure solver
-pressure_solver = SpectralPressureSolver(setup)
-
-# Time interval
-t_start, t_end = tlims = (T(0), T(1))
-
-# Initial conditions
-initial_velocity_u(x, y) = -sin(x)cos(y)
-initial_velocity_v(x, y) = cos(x)sin(y)
-initial_pressure(x, y) = 1 // 4 * (cos(2x) + cos(2y))
-V₀, p₀ = create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t_start;
- initial_pressure,
- pressure_solver,
-);
-
-# Solve steady state problem
-V, p = solve_steady_state(setup, V₀, p₀; npicard = 2);
-
-# Iteration processors
-processors = (
- ## field_plotter(setup; nupdate = 1),
- energy_history_plotter(setup; nupdate = 1),
- ## energy_spectrum_plotter(setup; nupdate = 1),
- ## animator(setup, "vorticity.mkv"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 1),
-);
-
-# Solve unsteady problem
-V, p, outputs = solve_unsteady(
- setup,
- V₀,
- p₀,
- tlims;
- Δt = T(0.01),
- processors,
- pressure_solver,
- inplace = true,
-);
-#md current_figure()
-
-# ## Post-process
-#
-# We may visualize or export the computed fields `(V, p)`
-
-# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
-
-# Plot pressure
-plot_pressure(setup, p)
-
-# Plot velocity
-plot_velocity(setup, V, t_end)
-
-# Plot vorticity
-plot_vorticity(setup, V, t_end)
-
-# Plot streamfunction
-## plot_streamfunction(setup, V, t_end)
-nothing
-
-# Energy history
-outputs[1]
+using LinearAlgebra
+
+"""
+Compare numerical solution with analytical solution at final time.
+"""
+function compute_convergence(; D, nlist, lims, Re, tlims, Δt, uref, ArrayType = Array)
+ T = typeof(lims[1])
+ e = zeros(T, length(nlist))
+ for (i, n) in enumerate(nlist)
+ @info "Computing error for n = $n"
+ x = ntuple(α -> LinRange(lims..., n + 1), D)
+ setup = Setup(x...; Re, ArrayType)
+ psolver = psolver_spectral(setup)
+ ustart = create_initial_conditions(
+ setup,
+ (dim, x...) -> uref(dim, x..., tlims[1]),
+ tlims[1];
+ psolver,
+ )
+ ut = create_initial_conditions(
+ setup,
+ (dim, x...) -> uref(dim, x..., tlims[2]),
+ tlims[2];
+ psolver,
+ doproject = false,
+ )
+ (; 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
+ e[i] = sqrt(a) / sqrt(b)
+ end
+ e
+end
+
+# Analytical solution for 2D Taylor-Green vortex
+solution(Re) =
+ (dim, x, y, t) -> (dim() == 1 ? -sin(x) * cos(y) : cos(x) * sin(y)) * exp(-2t / Re)
+
+# Compute error for different resolutions
+Re = 2.0e3
+nlist = [2, 4, 8, 16, 32, 64, 128, 256]
+e = compute_convergence(;
+ D = 2,
+ nlist,
+ lims = (0.0, 2π),
+ Re,
+ tlims = (0.0, 2.0),
+ Δt = 0.01,
+ uref = solution(Re),
+)
+
+# Plot convergence
+fig = Figure()
+ax = Axis(
+ fig[1, 1];
+ xscale = log10,
+ yscale = log10,
+ xticks = nlist,
+ xlabel = "n",
+ title = "Relative error",
+)
+scatterlines!(ax, nlist, e; label = "Data")
+lines!(ax, collect(extrema(nlist)), n -> n^-2.0; linestyle = :dash, label = "n^-2")
+axislegend(ax)
+fig
diff --git a/examples/TaylorGreenVortex3D.jl b/examples/TaylorGreenVortex3D.jl
index 41a200a9a..05ffd41e5 100644
--- a/examples/TaylorGreenVortex3D.jl
+++ b/examples/TaylorGreenVortex3D.jl
@@ -1,121 +1,82 @@
-# Little LSP hack to get function signatures, go #src
-# to definition etc. #src
-if isdefined(@__MODULE__, :LanguageServer) #src
- include("../src/IncompressibleNavierStokes.jl") #src
- using .IncompressibleNavierStokes #src
-end #src
-
# # Taylor-Green vortex - 3D
#
# In this example we consider the Taylor-Green vortex.
-# We start by loading packages.
-# A [Makie](https://github.com/JuliaPlots/Makie.jl) plotting backend is needed
-# for plotting. `GLMakie` creates an interactive window (useful for real-time
-# plotting), but does not work when building this example on GitHub.
-# `CairoMakie` makes high-quality static vector-graphics plots.
-
#md using CairoMakie
using GLMakie #!md
using IncompressibleNavierStokes
-# Case name for saving results
-name = "TaylorGreenVortex3D"
+# Output directory
+output = "output/TaylorGreenVortex3D"
# Floating point precision
-T = Float32
-
-# For CPU
-device = identity
+T = Float64
-# For GPU (note that `cu` converts to `Float32`)
-## using CUDA
-## device = cu
-nothing
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
-# Viscosity model
-viscosity_model = LaminarModel(; Re = T(2_000))
+# Reynolds number
+Re = T(6_000)
# A 3D grid is a Cartesian product of three vectors
n = 32
-lims = (T(0), T(2π))
+lims = T(0), T(1)
x = LinRange(lims..., n + 1)
y = LinRange(lims..., n + 1)
z = LinRange(lims..., n + 1)
-plot_grid(x, y, z)
# Build setup and assemble operators
-setup = Setup(x, y, z; viscosity_model);
+setup = Setup(x, y, z; Re, ArrayType);
# Since the grid is uniform and identical for x, y, and z, we may use a
# specialized spectral pressure solver
-pressure_solver = SpectralPressureSolver(setup)
+psolver = psolver_spectral(setup);
# Initial conditions
-initial_velocity_u(x, y, z) = sin(x)cos(y)cos(z)
-initial_velocity_v(x, y, z) = -cos(x)sin(y)cos(z)
-initial_velocity_w(x, y, z) = zero(x)
-initial_pressure(x, y, z) = 1 // 4 * (cos(2x) + cos(2y) + cos(2z))
-V₀, p₀ = create_initial_conditions(
+ustart = create_initial_conditions(
setup,
- initial_velocity_u,
- initial_velocity_v,
- initial_velocity_w,
- T(0);
- initial_pressure,
- pressure_solver,
-);
-
-# Solve steady state problem
-## V, p = solve_steady_state(setup, V₀, p₀; npicard = 6)
-nothing
-
-# Iteration processors
-processors = (
- ## field_plotter(device(setup); nupdate = 5),
- ## energy_history_plotter(device(setup); nupdate = 1),
- ## energy_spectrum_plotter(device(setup); nupdate = 100),
- ## animator(device(setup), "vorticity.mp4"; nupdate = 4),
- ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
- ## field_saver(setup; nupdate = 10),
- step_logger(; nupdate = 10),
+ (dim, x, y, z) ->
+ dim() == 1 ? sinpi(2x) * cospi(2y) * sinpi(2z) / 2 :
+ dim() == 2 ? -cospi(2x) * sinpi(2y) * sinpi(2z) / 2 : zero(x);
+ psolver,
);
-# Time interval
-t_start, t_end = tlims = (T(0), T(5))
-
# Solve unsteady problem
-V, p, outputs = solve_unsteady(
+(; u, t), outputs = solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- Δt = T(0.01),
- processors,
- pressure_solver,
- inplace = true,
- device,
+ ustart,
+ tlims = (T(0), T(1.0)),
+ Δt = T(1e-3),
+ processors = (
+ ## rtp = realtimeplotter(; setup, plot = fieldplot, nupdate = 10),
+ ehist = realtimeplotter(;
+ setup,
+ plot = energy_history_plot,
+ nupdate = 1,
+ displayfig = false,
+ ),
+ espec = realtimeplotter(; setup, plot = energy_spectrum_plot, nupdate = 10),
+ ## anim = animator(; setup, path = "$output/solution.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 100),
+ ),
+ psolver,
);
# ## Post-process
#
-# We may visualize or export the computed fields `(V, p)`
-
-# Export to VTK
-save_vtk(setup, V, p, t_end, "output/solution")
+# We may visualize or export the computed fields
-# Plot pressure
-nothing #md
-plot_pressure(setup, p; levels = 3, alpha = 0.05) #!md
+# Energy history
+outputs.ehist
-# Plot velocity
-nothing #md
-plot_velocity(setup, V, t_end; levels = 3, alpha = 0.05) #!md
+# Energy spectrum
+outputs.espec
-# Plot vorticity
-nothing #md
-plot_vorticity(setup, V, t_end; levels = 5, alpha = 0.05) #!md
-
-# Plot streamfunction
-## plot_streamfunction(setup, V, t_end)
-nothing
+# Export to VTK
+save_vtk(setup, u, t, "$output/solution"; psolver)
diff --git a/examples/src/Examples.jl b/examples/src/Examples.jl
new file mode 100644
index 000000000..7d9f903f1
--- /dev/null
+++ b/examples/src/Examples.jl
@@ -0,0 +1,5 @@
+"""
+Empty module for examples environment.
+This allows the docs to load example dependencies.
+"""
+module Examples end
diff --git a/ext/IncompressibleNavierStokesCUDAExt.jl b/ext/IncompressibleNavierStokesCUDAExt.jl
new file mode 100644
index 000000000..38f91edb1
--- /dev/null
+++ b/ext/IncompressibleNavierStokesCUDAExt.jl
@@ -0,0 +1,56 @@
+"""
+ IncompressibleNavierStokesCUDAExt
+
+CUDA extension for IncompressibleNavierStokes.
+"""
+module IncompressibleNavierStokesCUDAExt
+
+using CUDA
+using CUDA.CUSPARSE
+using CUDSS
+using IncompressibleNavierStokes
+using IncompressibleNavierStokes: PressureBC, laplacian_mat
+using SparseArrays
+
+# CUDA version, using CUDSS LDLt decomposition.
+function IncompressibleNavierStokes.psolver_direct(::CuArray, setup)
+ (; grid, boundary_conditions) = setup
+ (; x, Np, Ip) = grid
+ T = eltype(x[1])
+ L = laplacian_mat(setup)
+ isdefinite =
+ any(bc -> bc[1] isa PressureBC || bc[2] isa PressureBC, boundary_conditions)
+ if isdefinite
+ # No extra DOF
+ ftemp = fill!(similar(x[1], prod(Np)), 0)
+ ptemp = fill!(similar(x[1], prod(Np)), 0)
+ viewrange = (:)
+ # structure = "SPD" # Symmetric positive definite
+ structure = "G" # General matrix
+ _view = 'F' # Full matrix representation
+ else
+ # With extra DOF
+ ftemp = fill!(similar(x[1], prod(Np) + 1), 0)
+ ptemp = fill!(similar(x[1], prod(Np) + 1), 0)
+ # e = fill!(similar(x[1], prod(Np)), 1)
+ e = ones(T, prod(Np))
+ L = SparseMatrixCSC(L)
+ L = [L e; e' 0]
+ viewrange = 1:prod(Np)
+ structure = "S" # Symmetric (not positive definite)
+ _view = 'L' # Lower triangular representation
+ end
+ L = CuSparseMatrixCSR(L)
+ solver = CudssSolver(L, structure, _view)
+ cudss("analysis", solver, ptemp, ftemp)
+ cudss("factorization", solver, ptemp, ftemp) # Compute factorization
+ function psolve!(p, f)
+ T = eltype(p)
+ copyto!(view(ftemp, viewrange), view(view(f, Ip), :))
+ cudss("solve", solver, ptemp, ftemp)
+ copyto!(view(view(p, Ip), :), eltype(p).(view(ptemp, viewrange)))
+ p
+ end
+end
+
+end
diff --git a/lib/NeuralClosure/LICENSE b/lib/NeuralClosure/LICENSE
new file mode 100644
index 000000000..d76a52488
--- /dev/null
+++ b/lib/NeuralClosure/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Syver Døving Agdestein, Benjamin Sanderse, and contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/lib/NeuralClosure/Project.toml b/lib/NeuralClosure/Project.toml
new file mode 100644
index 000000000..138073932
--- /dev/null
+++ b/lib/NeuralClosure/Project.toml
@@ -0,0 +1,30 @@
+name = "NeuralClosure"
+uuid = "099dac27-d7f2-4047-93d5-0baee36b9c25"
+authors = ["Syver Døving Agdestein "]
+version = "0.1.0"
+
+[deps]
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
+IncompressibleNavierStokes = "5e318141-6589-402b-868d-77d7df8c442e"
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+Observables = "510215fc-4207-5dde-b226-833fc4488ee2"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
+
+[compat]
+CUDA = "5"
+ComponentArrays = "0.15"
+KernelAbstractions = "0.9"
+Lux = "0.5"
+NNlib = "0.9"
+Observables = "0.5"
+Optimisers = "0.3"
+Random = "1"
+Tullio = "0.3"
+Zygote = "0.6"
+julia = "1.7"
diff --git a/lib/NeuralClosure/README.md b/lib/NeuralClosure/README.md
new file mode 100644
index 000000000..c3f08b77b
--- /dev/null
+++ b/lib/NeuralClosure/README.md
@@ -0,0 +1,25 @@
+## NeuralClosure
+
+Neural closure modeling tools for
+[IncompressibleNavierStokes.jl](https://github.com/agdestein/IncompressibleNavierStokes.jl).
+
+## Set up environment
+
+From this directory, run:
+
+```sh
+julia --project -e '
+using Pkg
+Pkg.develop(PackageSpec(; path = "../.."))
+Pkg.instantiate()
+'
+```
+
+or interactively from a Julia REPL:
+
+```julia-repl
+julia> ]
+(v1.10) pkg> activate .
+(NeuralClosure) pkg> dev ../..
+(NeuralClosure) pkg> instantiate
+```
diff --git a/lib/NeuralClosure/src/NeuralClosure.jl b/lib/NeuralClosure/src/NeuralClosure.jl
new file mode 100644
index 000000000..18ae25a3f
--- /dev/null
+++ b/lib/NeuralClosure/src/NeuralClosure.jl
@@ -0,0 +1,37 @@
+"""
+Neural closure modelling tools.
+"""
+module NeuralClosure
+
+using ComponentArrays: ComponentArray
+using IncompressibleNavierStokes
+using IncompressibleNavierStokes: Dimension, momentum!, apply_bc_u!, project!
+using KernelAbstractions
+using Lux
+using NNlib
+using Observables
+using Random
+using Tullio
+using Zygote
+
+# Must be loaded inside for Tullio to work correctly
+using CUDA
+
+include("closure.jl")
+include("cnn.jl")
+include("fno.jl")
+include("training.jl")
+include("filter.jl")
+include("create_les_data.jl")
+
+export smagorinsky_closure
+export cnn, fno, FourierLayer
+export train
+export mean_squared_error, create_relerr_prior, create_relerr_post
+export create_loss_prior, create_loss_post
+export create_dataloader_prior, create_dataloader_post
+export create_callback, create_les_data, create_io_arrays
+export wrappedclosure
+export FaceAverage, VolumeAverage, reconstruct, reconstruct!
+
+end
diff --git a/lib/NeuralClosure/src/closure.jl b/lib/NeuralClosure/src/closure.jl
new file mode 100644
index 000000000..b2faa3499
--- /dev/null
+++ b/lib/NeuralClosure/src/closure.jl
@@ -0,0 +1,137 @@
+"""
+ wrappedclosure(m, setup)
+
+Wrap closure model and parameters so that it can be used in the solver.
+"""
+function wrappedclosure(m, setup)
+ (; dimension, Iu) = setup.grid
+ D = dimension()
+ # function neuralclosure(u)
+ # u = stack(ntuple(α -> u[α][Iu[α]], D))
+ # u = reshape(u, size(u)..., 1) # One sample
+ # mu = m(u, θ)
+ # mu = pad_circular(mu, 1)
+ # sz..., _ = size(mu)
+ # i = ntuple(Returns(:), D)
+ # mu = ntuple(α -> mu[i..., α, 1], D)
+ # end
+ neuralclosure(u, θ) =
+ if D == 2
+ u = cat(u[1][Iu[1]], u[2][Iu[2]]; dims = 3)
+ u = reshape(u, size(u)..., 1) # One sample
+ mu = m(u, θ)
+ mu = pad_circular(mu, 1)
+ mu = (mu[:, :, 1, 1], mu[:, :, 2, 1])
+ elseif D == 3
+ u = cat(u[1][Iu[1]], u[2][Iu[2]], u[3][Iu[3]]; dims = 4)
+ u = reshape(u, size(u)..., 1) # One sample
+ mu = m(u, θ)
+ mu = pad_circular(mu, 1)
+ mu = (mu[:, :, :, 1, 1], mu[:, :, :, 2, 1], mu[:, :, :, 3, 1])
+ end
+end
+
+"""
+ create_closure(layers...; rng)
+
+Create neural closure model from layers.
+"""
+function create_closure(layers...; rng)
+ chain = Chain(layers...)
+
+ # Create parameter vector (empty state)
+ params, state = Lux.setup(rng, chain)
+ θ = ComponentArray(params)
+
+ # Compute closure term for given parameters
+ closure(u, θ) = first(chain(u, θ, state))
+
+ closure, θ
+end
+
+"""
+ create_tensorclosure(layers...; setup, rng)
+
+Create tensor basis closure.
+"""
+function create_tensorclosure(layers...; setup, rng)
+ D = setup.grid.dimension()
+ cnn, θ = create_closure(layers...; rng)
+ function closure(u, θ)
+ B, V = tensorbasis(u, setup)
+ V = stack(V)
+ α = cnn(V, θ)
+ τ = sum(k -> α[ntuple(Returns(:), D)..., k] .* B[k], 1:length(B))
+ end
+end
+
+"""
+ collocate(u)
+
+Interpolate velocity components to volume centers.
+"""
+function collocate(u)
+ sz..., D, _ = size(u)
+ # for α = 1:D
+ # v = selectdim(u, D + 1, α)
+ # v = (v + circshift(v, ntuple(β -> α == β ? -1 : 0, D + 1))) / 2
+ # end
+ if D == 2
+ a = selectdim(u, 3, 1)
+ b = selectdim(u, 3, 2)
+ a = (a .+ circshift(a, (-1, 0, 0))) ./ 2
+ b = (b .+ circshift(b, (0, -1, 0))) ./ 2
+ a = reshape(a, sz..., 1, :)
+ b = reshape(b, sz..., 1, :)
+ cat(a, b; dims = 3)
+ elseif D == 3
+ a = selectdim(u, 4, 1)
+ b = selectdim(u, 4, 2)
+ c = selectdim(u, 4, 3)
+ a = (a .+ circshift(a, (-1, 0, 0, 0))) ./ 2
+ b = (b .+ circshift(b, (0, -1, 0, 0))) ./ 2
+ c = (c .+ circshift(c, (0, 0, -1, 0))) ./ 2
+ a = reshape(a, sz..., 1, :)
+ b = reshape(b, sz..., 1, :)
+ c = reshape(c, sz..., 1, :)
+ cat(a, b, c; dims = 4)
+ end
+end
+
+"""
+ decollocate(u)
+
+Interpolate closure force from volume centers to volume faces.
+"""
+function decollocate(u)
+ sz..., D, _ = size(u)
+ # for α = 1:D
+ # v = selectdim(u, D + 1, α)
+ # v = (v + circshift(v, ntuple(β -> α == β ? -1 : 0, D + 1))) / 2
+ # end
+ if D == 2
+ a = selectdim(u, 3, 1)
+ b = selectdim(u, 3, 2)
+ a = (a .+ circshift(a, (1, 0, 0))) ./ 2
+ b = (b .+ circshift(b, (0, 1, 0))) ./ 2
+ # a = circshift(a, (1, 0, 0)) .- a
+ # b = circshift(b, (0, 1, 0)) .- b
+ a = reshape(a, sz..., 1, :)
+ b = reshape(b, sz..., 1, :)
+ cat(a, b; dims = 3)
+ elseif D == 3
+ a = selectdim(u, 4, 1)
+ b = selectdim(u, 4, 2)
+ c = selectdim(u, 4, 3)
+ a = (a .+ circshift(a, (1, 0, 0, 0))) ./ 2
+ b = (b .+ circshift(b, (0, 1, 0, 0))) ./ 2
+ c = (c .+ circshift(c, (0, 0, 1, 0))) ./ 2
+ # a = circshift(a, (1, 0, 0, 0)) .- a
+ # b = circshift(b, (0, 1, 0, 0)) .- b
+ # c = circshift(c, (0, 0, 1, 0)) .- c
+ a = reshape(a, sz..., 1, :)
+ b = reshape(b, sz..., 1, :)
+ c = reshape(c, sz..., 1, :)
+ cat(a, b, c; dims = 4)
+ end
+end
diff --git a/lib/NeuralClosure/src/cnn.jl b/lib/NeuralClosure/src/cnn.jl
new file mode 100644
index 000000000..53ed71353
--- /dev/null
+++ b/lib/NeuralClosure/src/cnn.jl
@@ -0,0 +1,70 @@
+"""
+ cnn(;
+ setup,
+ radii,
+ channels,
+ activations,
+ use_bias,
+ channel_augmenter = identity,
+ rng = Random.default_rng(),
+ )
+
+Create CNN closure model. Return a tuple `(closure, θ)` where `θ` are the initial
+parameters and `closure(u, θ)` predicts the commutator error.
+"""
+function cnn(;
+ setup,
+ radii,
+ channels,
+ activations,
+ use_bias,
+ channel_augmenter = identity,
+ rng = Random.default_rng(),
+)
+ r, c, σ, b = radii, channels, activations, use_bias
+ (; T, grid, boundary_conditions) = setup
+ (; dimension) = grid
+ D = dimension()
+
+ # dx = map(d -> d[2:end-1], Δu)
+
+ # Weight initializer
+ glorot_uniform_T(rng::AbstractRNG, dims...) = glorot_uniform(rng, T, dims...)
+
+ # Make sure there are two force fields in output
+ @assert c[end] == D
+
+ # Add input channel size
+ c = [D; c]
+
+ # Create convolutional closure model
+ layers = (
+ # Put inputs in pressure points
+ collocate,
+
+ # Add padding so that output has same shape as commutator error
+ ntuple(
+ α ->
+ boundary_conditions[α][1] isa PeriodicBC ?
+ u -> pad_circular(u, sum(r); dims = α) :
+ u -> pad_repeat(u, sum(r); dims = α),
+ D,
+ ),
+
+ # Some convolutional layers
+ (
+ Conv(
+ ntuple(α -> 2r[i] + 1, D),
+ c[i] => c[i+1],
+ σ[i];
+ use_bias = b[i],
+ init_weight = glorot_uniform_T,
+ ) for i ∈ eachindex(r)
+ )...,
+
+ # Differentiate output to velocity points
+ decollocate,
+ )
+
+ create_closure(layers...; rng)
+end
diff --git a/lib/NeuralClosure/src/create_les_data.jl b/lib/NeuralClosure/src/create_les_data.jl
new file mode 100644
index 000000000..93125a3e4
--- /dev/null
+++ b/lib/NeuralClosure/src/create_les_data.jl
@@ -0,0 +1,252 @@
+# Body force
+function gaussian_force(
+ x,
+ y;
+ σ = eltype(x)(0.05),
+ A = eltype(x)(0.002),
+ rng = Random.default_rng(),
+)
+ T = eltype(x)
+ Lx = x[end] - x[1]
+ Ly = y[end] - y[1]
+ xc = x[1] + rand(rng, T) * Lx
+ yc = y[1] + rand(rng, T) * Ly
+ σx = σ * Lx
+ σy = σ * Ly
+ ϕ = T(2π) * rand(rng, T)
+ f = sum(
+ [
+ A * exp(-(x - xc - lx)^2 / 2σx^2 - (y - yc - ly)^2 / 2σy^2) for
+ y ∈ y[2:end], x ∈ x[2:end], ly ∈ (-Ly, T(0), Ly), lx ∈ (-Lx, T(0), Lx)
+ ];
+ dims = (3, 4),
+ )[
+ :,
+ :,
+ 1,
+ 1,
+ ]
+ force = cat(sin(ϕ) * f, cos(ϕ) * f; dims = 3)
+ force = reshape(force, :)
+ force = force .- sum(force) / length(force)
+ force
+end
+
+function lesdatagen(dnsobs, Φ, les, compression, psolver)
+ Φu = zero.(Φ(dnsobs[].u, les, compression))
+ p = zero(Φu[1])
+ div = zero(p)
+ ΦF = zero.(Φu)
+ FΦ = zero.(Φu)
+ c = zero.(Φu)
+ results = (; u = fill(Array.(dnsobs[].u), 0), c = fill(Array.(dnsobs[].u), 0))
+ temp = nothing
+ on(dnsobs) do (; u, F, t)
+ Φ(Φu, u, les, compression)
+ apply_bc_u!(Φu, t, les)
+ Φ(ΦF, F, les, compression)
+ momentum!(FΦ, Φu, temp, t, les)
+ apply_bc_u!(FΦ, t, les; dudt = true)
+ project!(FΦ, les; psolver, div, p)
+ for α = 1:length(u)
+ c[α] .= ΦF[α] .- FΦ[α]
+ end
+ push!(results.u, Array.(Φu))
+ push!(results.c, Array.(c))
+ end
+ results
+end
+
+"""
+ filtersaver(dns, les, filters, compression, psolver_dns, psolver_les; nupdate = 1)
+
+Save filtered DNS data.
+"""
+filtersaver(dns, les, filters, compression, psolver_dns, psolver_les; nupdate = 1) =
+ processor(
+ (results, state) -> (; results..., comptime = time() - results.comptime),
+ ) do state
+ comptime = time()
+ (; x) = dns.grid
+ T = eltype(x[1])
+ F = zero.(state[].u)
+ div = zero(state[].u[1])
+ p = zero(state[].u[1])
+ dnsobs = Observable((; state[].u, F, state[].t))
+ data = [
+ lesdatagen(dnsobs, Φ, les[i], compression[i], psolver_les[i]) for
+ i = 1:length(les), Φ in filters
+ ]
+ results = (; data, t = zeros(T, 0), comptime)
+ temp = nothing
+ on(state) do (; u, t, n)
+ n % nupdate == 0 || return
+ momentum!(F, u, temp, t, dns)
+ apply_bc_u!(F, t, dns; dudt = true)
+ project!(F, dns; psolver = psolver_dns, div, p)
+ push!(results.t, t)
+ dnsobs[] = (; u, F, t)
+ end
+ state[] = state[] # Save initial conditions
+ results
+ end
+
+"""
+ create_les_data(
+ D = 2,
+ Re = 2e3,
+ lims = ntuple(α -> (typeof(Re)(0), typeof(Re)(1)), D),
+ nles = [ntuple(α -> 64, D)],
+ ndns = ntuple(α -> 256, D),
+ filters = (FaceAverage(),),
+ tburn = typeof(Re)(0.1),
+ tsim = typeof(Re)(0.1),
+ Δt = typeof(Re)(1e-4),
+ create_psolver = psolver_spectral,
+ savefreq = 1,
+ ArrayType = Array,
+ icfunc = (setup, psolver) -> random_field(setup, typeof(Re)(0); psolver),
+ rng,
+ kwargs...,
+ )
+
+Create filtered DNS data.
+"""
+function create_les_data(;
+ D = 2,
+ Re = 2e3,
+ lims = ntuple(α -> (typeof(Re)(0), typeof(Re)(1)), D),
+ nles = [ntuple(α -> 64, D)],
+ ndns = ntuple(α -> 256, D),
+ filters = (FaceAverage(),),
+ tburn = typeof(Re)(0.1),
+ tsim = typeof(Re)(0.1),
+ Δt = typeof(Re)(1e-4),
+ create_psolver = psolver_spectral,
+ savefreq = 1,
+ ArrayType = Array,
+ icfunc = (setup, psolver, rng) -> random_field(setup, typeof(Re)(0); psolver, rng),
+ rng,
+ kwargs...,
+)
+ T = typeof(Re)
+
+ compression = [ndns[1] ÷ nles[1] for nles in nles]
+ for (c, n) in zip(compression, nles), α = 1:D
+ @assert c * n[α] == ndns[α]
+ end
+
+ # Build setup and assemble operators
+ dns = Setup(
+ ntuple(α -> LinRange(lims[α]..., ndns[α] + 1), D)...;
+ Re,
+ ArrayType,
+ kwargs...,
+ )
+ les = [
+ Setup(
+ ntuple(α -> LinRange(lims[α]..., nles[α] + 1), D)...;
+ Re,
+ ArrayType,
+ kwargs...,
+ ) for nles in nles
+ ]
+
+ # Since the grid is uniform and identical for x and y, we may use a specialized
+ # spectral pressure solver
+ psolver = create_psolver(dns)
+ psolver_les = create_psolver.(les)
+
+ # Number of time steps to save
+ nt = round(Int, tsim / Δt)
+ Δt = tsim / nt
+
+ # datasize = Base.summarysize(filtered) / 1e6
+ datasize =
+ length(filters) *
+ (nt ÷ savefreq + 1) *
+ sum(prod.(nles)) *
+ D *
+ 2 *
+ length(bitstring(zero(T))) / 8 / 1e6
+ @info "Generating $datasize Mb of filtered DNS data"
+
+ # Initial conditions
+ u₀ = icfunc(dns, psolver, rng)
+
+ any(u -> any(isnan, u), u₀) && @warn "Initial conditions contain NaNs"
+
+ # Random body force
+ # force_dns =
+ # gaussian_force(xdns...) +
+ # gaussian_force(xdns...) +
+ # # gaussian_force(xdns...) +
+ # # gaussian_force(xdns...) +
+ # gaussian_force(xdns...)
+ # force_dns = zero.(u₀)
+ # force_les = face_average(force_dns, les, compression)
+
+ _dns = dns
+ _les = les
+ # _dns = (; dns..., bodyforce = force_dns)
+ # _les = (; les..., bodyforce = force_les)
+
+ # Solve burn-in DNS
+ (; u, t), outputs = solve_unsteady(_dns, u₀, (T(0), tburn); Δt, psolver)
+
+ # Solve DNS and store filtered quantities
+ (; u, t), outputs = solve_unsteady(
+ _dns,
+ u,
+ (T(0), tsim);
+ Δt,
+ processors = (;
+ f = filtersaver(
+ _dns,
+ _les,
+ filters,
+ compression,
+ psolver,
+ psolver_les;
+ nupdate = savefreq,
+ ),
+ # plot = realtimeplotter(; setup = dns, nupdate = 10),
+ log = timelogger(; nupdate = 10),
+ ),
+ psolver,
+ )
+
+ # Store result for current IC
+ outputs.f
+end
+
+"""
+ create_io_arrays(data, setups)
+
+Create ``(\\bar{u}, c)`` pairs for training.
+"""
+function create_io_arrays(data, setups)
+ nsample = length(data)
+ ngrid, nfilter = size(data[1].data)
+ nt = length(data[1].t) - 1
+ T = eltype(data[1].t)
+ map(CartesianIndices((ngrid, nfilter))) do I
+ ig, ifil = I.I
+ (; dimension, N, Iu) = setups[ig].grid
+ D = dimension()
+ u = zeros(T, (N .- 2)..., D, nt + 1, nsample)
+ c = zeros(T, (N .- 2)..., D, nt + 1, nsample)
+ ifield = ntuple(Returns(:), D)
+ for is = 1:nsample, it = 1:nt+1, α = 1:D
+ copyto!(
+ view(u, ifield..., α, it, is),
+ view(data[is].data[ig, ifil].u[it][α], Iu[α]),
+ )
+ copyto!(
+ view(c, ifield..., α, it, is),
+ view(data[is].data[ig, ifil].c[it][α], Iu[α]),
+ )
+ end
+ (; u = reshape(u, (N .- 2)..., D, :), c = reshape(c, (N .- 2)..., D, :))
+ end
+end
diff --git a/lib/NeuralClosure/src/filter.jl b/lib/NeuralClosure/src/filter.jl
new file mode 100644
index 000000000..dfb1c4943
--- /dev/null
+++ b/lib/NeuralClosure/src/filter.jl
@@ -0,0 +1,122 @@
+abstract type AbstractFilter end
+
+struct FaceAverage <: AbstractFilter end
+struct VolumeAverage <: AbstractFilter end
+
+(Φ::AbstractFilter)(u, setup_les, compression) = Φ(
+ ntuple(α -> fill!(similar(u[1], setup_les.grid.N), 0), length(u)),
+ u,
+ setup_les,
+ compression,
+)
+
+"""
+ (::FaceAverage)(v, u, setup_les)
+
+Average fine grid `u` over coarse volume face. Put result in `v`.
+"""
+function (::FaceAverage)(v, u, setup_les, comp)
+ (; grid, workgroupsize) = setup_les
+ (; Nu, Iu) = grid
+ D = length(u)
+ @kernel function Φ!(v, u, ::Val{α}, face, I0) where {α}
+ I = @index(Global, Cartesian)
+ J = I0 + comp * (I - oneunit(I))
+ s = zero(eltype(v[α]))
+ for i in face
+ s += u[α][J+i]
+ end
+ v[α][I0+I] = s / comp^(D - 1)
+ end
+ for α = 1:D
+ ndrange = Nu[α]
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ face = CartesianIndices(ntuple(β -> β == α ? (comp:comp) : (1:comp), D))
+ Φ!(get_backend(v[1]), workgroupsize)(v, u, Val(α), face, I0; ndrange)
+ end
+ v
+end
+
+"""
+ reconstruct!(u, v, setup_dns, setup_les)
+
+Average fine grid `u` over coarse volume face. Put result in `v`.
+"""
+function reconstruct!(u, v, setup_dns, setup_les, comp)
+ (; grid, boundary_conditions, workgroupsize) = setup_les
+ (; N, Iu) = grid
+ D = length(u)
+ e = Offset{D}()
+ @assert all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions)
+ @kernel function R!(u, v, ::Val{α}, volume) where {α}
+ J = @index(Global, Cartesian)
+ I = oneunit(J) + comp * J
+ J = oneunit(J) + J
+ Jleft = J - e(α)
+ Jleft.I[α] == 1 && (Jleft += (N[α] - 2) * e(α))
+ for i in volume
+ s = zero(eltype(v[α]))
+ s += (comp - i.I[α]) * v[α][J]
+ s += i.I[α] * v[α][Jleft]
+ u[α][I-i] = s / comp
+ end
+ end
+ for α = 1:D
+ ndrange = N .- 2
+ volume = CartesianIndices(ntuple(β -> 0:comp-1, D))
+ R!(get_backend(v[1]), workgroupsize)(u, v, Val(α), volume; ndrange)
+ end
+ u
+end
+
+reconstruct(v, setup_dns, setup_les, comp) = reconstruct!(
+ ntuple(α -> fill!(similar(v[1], setup_dns.grid.N), 0), length(v)),
+ v,
+ setup_dns,
+ setup_les,
+ comp,
+)
+
+"""
+ (::VolumeAverage)(v, u, setup_les, comp)
+
+Average fine grid `u` over coarse volume. Put result in `v`.
+"""
+function (::VolumeAverage)(v, u, setup_les, comp)
+ (; grid, boundary_conditions, workgroupsize) = setup_les
+ (; N, Nu, Iu) = grid
+ D = length(u)
+ @assert all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions)
+ @kernel function Φ!(v, u, ::Val{α}, volume, I0) where {α}
+ I = @index(Global, Cartesian)
+ J = I0 + comp * (I - oneunit(I))
+ s = zero(eltype(v[α]))
+ # n = 0
+ for i in volume
+ # Periodic extension
+ K = J + i
+ K = mod1.(K.I, comp .* (N .- 2))
+ K = CartesianIndex(K)
+ s += u[α][K]
+ # n += 1
+ end
+ n = (iseven(comp) ? comp + 1 : comp) * comp^(D - 1)
+ v[α][I0+I] = s / n
+ end
+ for α = 1:D
+ ndrange = Nu[α]
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ volume = CartesianIndices(
+ ntuple(
+ β ->
+ α == β ? iseven(comp) ? (comp÷2:comp÷2+comp) : (comp+1:2comp+1) :
+ (1:comp),
+ D,
+ ),
+ )
+ Φ!(get_backend(v[1]), workgroupsize)(v, u, Val(α), volume, I0; ndrange)
+ end
+ v
+end
diff --git a/lib/NeuralClosure/src/fno.jl b/lib/NeuralClosure/src/fno.jl
new file mode 100644
index 000000000..38edb78da
--- /dev/null
+++ b/lib/NeuralClosure/src/fno.jl
@@ -0,0 +1,210 @@
+"""
+ fno(; setup, kmax, c, σ, ψ, rng = Random.default_rng(), kwargs...)
+
+Create FNO closure model. Return a tuple `(closure, θ)` where `θ` are the
+initial parameters and `closure(V, θ)` predicts the commutator error.
+"""
+function fno(; setup, kmax, c, σ, ψ, rng = Random.default_rng(), kwargs...)
+ (; grid) = setup
+ (; dimension, x, N) = grid
+
+ D = dimension()
+
+ @assert all(==(first(N)), N)
+
+ # Fourier layers
+ @assert length(kmax) == length(c) == length(σ)
+
+ # Make sure there are two velocity fields in input and output
+ c = [D; c]
+
+ # Weight initializer
+ T = eltype(x[1])
+ init_weight(rng::AbstractRNG, dims...) = glorot_uniform(rng, T, dims...)
+
+ # Create FNO closure model
+ layers = (
+ # Put inputs in pressure points
+ collocate,
+
+ # Conv(ntuple(Returns(1), D), D => c[1]; use_bias = false, init_weight),
+
+ # Some Fourier layers
+ (
+ FourierLayer(dimension, kmax[i], c[i] => c[i+1]; σ = σ[i], init_weight) for
+ i ∈ eachindex(σ)
+ )...,
+
+ # Compress with a final dense layer
+ # Conv(ntuple(Returns(1), D), c[end] => D; use_bias = false, init_weight),
+ Conv(ntuple(Returns(1), D), c[end] => 2 * c[end], ψ; init_weight),
+ Conv(ntuple(Returns(1), D), 2 * c[end] => D; use_bias = false, init_weight),
+
+ # Differentiate output to velocity points
+ decollocate,
+ )
+ create_closure(layers...; rng)
+end
+
+"""
+ FourierLayer(dimension, kmax, cin => cout; σ = identity, init_weight = glorot_uniform)
+
+Fourier layer operating on uniformly discretized functions.
+
+Some important sizes:
+
+- `dimension`: Spatial dimension, e.g. `Dimension(2)` or `Dimension(3)`.
+- `(nx..., cin, nsample)`: Input size
+- `(nx..., cout, nsample)`: Output size
+- `nx = fill(n, dimension())`: Number of points in each spatial dimension
+- `n ≥ kmax`: Same number of points in each spatial dimension, must be
+ larger than cut-off wavenumber
+- `kmax`: Cut-off wavenumber
+- `nsample`: Number of input samples (treated independently)
+"""
+struct FourierLayer{D,A,F} <: Lux.AbstractExplicitLayer
+ dimension::Dimension{D}
+ kmax::Int
+ cin::Int
+ cout::Int
+ σ::A
+ init_weight::F
+end
+
+FourierLayer(
+ dimension,
+ kmax,
+ ch::Pair{Int,Int};
+ σ = identity,
+ init_weight = glorot_uniform,
+) = FourierLayer(dimension, kmax, first(ch), last(ch), σ, init_weight)
+
+Lux.initialparameters(
+ rng::AbstractRNG,
+ (; dimension, kmax, cin, cout, init_weight)::FourierLayer,
+) = (;
+ spatial_weight = init_weight(rng, cout, cin),
+ # spectral_weights = init_weight(rng, fill(kmax + 1, dimension())..., cout, cin, 2),
+ spectral_weights = init_weight(rng, fill(2 * (kmax + 1), dimension())..., cout, cin, 2),
+)
+Lux.initialstates(::AbstractRNG, ::FourierLayer) = (;)
+# Lux.parameterlength((; dimension, kmax, cin, cout)::FourierLayer) =
+# cout * cin + (kmax + 1)^dimension() * 2 * cout * cin
+Lux.parameterlength((; dimension, kmax, cin, cout)::FourierLayer) =
+ cout * cin + (2 * (kmax + 1))^dimension() * 2 * cout * cin
+Lux.statelength(::FourierLayer) = 0
+
+## Pretty printing
+function Base.show(io::IO, (; dimension, kmax, cin, cout, σ)::FourierLayer)
+ print(io, "FourierLayer{", dimension(), "}(")
+ print(io, kmax)
+ print(io, ", ", cin, " => ", cout)
+ print(io, "; σ = ", σ)
+ print(io, ")")
+end
+
+# Pass inputs through Fourier layer
+function ((; dimension, kmax, cout, cin, σ)::FourierLayer)(x, params, state)
+ # TODO: Check if this is more efficient for
+ # size(x) = (cin, nx..., nsample) or
+ # size(x) = (nx..., cin, nsample)
+
+ # TODO: Set FFT normalization so that layer is truly grid independent
+
+ # Spatial dimension
+ D = dimension()
+
+ nx..., _cin, nsample = size(x)
+ K = nx[1]
+ @assert _cin == cin "Number of input channels must be compatible with weights"
+ @assert all(==(first(nx)), nx) "Fourier layer requires same number of grid points in each dimension"
+ @assert kmax ≤ first(nx) "Fourier layer input must be discretized on at least `kmax` points"
+
+ # Destructure params
+ # The real and imaginary parts of R are stored in two separate channels
+ W = params.spatial_weight
+ R = params.spectral_weights
+ R = selectdim(R, D + 3, 1) .+ im .* selectdim(R, D + 3, 2)
+
+ # Spatial part (applied point-wise)
+ if D == 2
+ @tullio y[i₁, i₂, b, s] := W[b, a] * x[i₁, i₂, a, s]
+ elseif D == 3
+ @tullio y[i₁, i₂, i₃, b, s] := W[b, a] * x[i₁, i₂, i₃, a, s]
+ end
+
+ # Spectral part (applied mode-wise)
+ #
+ # Steps:
+ #
+ # - go to complex-valued spectral space
+ # - chop off high wavenumbers
+ # - multiply with weights mode-wise
+ # - pad with zeros to restore original shape
+ # - go back to real valued spatial representation
+ ikeep = ntuple(Returns([1:kmax+1; K-kmax:K]), D)
+ # ikeep = ntuple(Returns(1:kmax+1), D)
+ dims = ntuple(identity, D)
+ xhat = fft(x, dims)
+ xhat = xhat[ikeep..., :, :]
+ if D == 2
+ @tullio z[k₁, k₂, b, s] := R[k₁, k₂, b, a] * xhat[k₁, k₂, a, s]
+ z = cat(
+ z[1:kmax+1, :, :, :],
+ zero(similar(z, K - 2 * (kmax + 1), 2 * (kmax + 1), cout, nsample)),
+ z[end-kmax:end, :, :, :];
+ dims = 1,
+ )
+ z = cat(
+ z[:, 1:kmax+1, :, :],
+ zero(similar(z, K, K - 2 * (kmax + 1), cout, nsample)),
+ z[:, end-kmax:end, :, :];
+ dims = 2,
+ )
+ elseif D == 3
+ @tullio z[k₁, k₂, k₃, b, s] := R[k₁, k₂, k₃, b, a] * xhat[k₁, k₂, k₃, a, s]
+ z = cat(
+ z[1:kmax+1, :, :, :, :],
+ zero(
+ similar(
+ z,
+ K - 2 * (kmax + 1),
+ 2 * (kmax + 1),
+ 2 * (kmax + 1),
+ cout,
+ nsample,
+ ),
+ ),
+ z[end-kmax:end, :, :, :, :];
+ dims = 1,
+ )
+ z = cat(
+ z[:, 1:kmax+1, :, :, :],
+ zero(similar(z, K, K - 2 * (kmax + 1), 2 * (kmax + 1), cout, nsample)),
+ z[:, end-kmax:end, :, :, :];
+ dims = 2,
+ )
+ z = cat(
+ z[:, :, 1:kmax+1, :, :],
+ zero(similar(z, K, K, K - 2 * (kmax + 1), cout, nsample)),
+ z[:, :, end-kmax:end, :, :];
+ dims = 3,
+ )
+ end
+ # z = pad_zeros(z, ntuple(i -> isodd(i) ? 0 : first(nx) - kmax - 1, 2 * D); dims)
+ z = real.(ifft(z, dims))
+
+ # Outer layer: Activation over combined spatial and spectral parts
+ # Note: Even though high wavenumbers are chopped off in `z` and may
+ # possibly not be present in the input at all, `σ` creates new high
+ # wavenumbers. High wavenumber functions may thus be represented using a
+ # sequence of Fourier layers. In this case, the `y`s are the only place
+ # where information contained in high input wavenumbers survive in a
+ # Fourier layer.
+ v = σ.(y .+ z)
+
+ # @infiltrate
+
+ # Fourier layer does not modify state
+ v, state
+end
diff --git a/lib/NeuralClosure/src/training.jl b/lib/NeuralClosure/src/training.jl
new file mode 100644
index 000000000..a663bf0df
--- /dev/null
+++ b/lib/NeuralClosure/src/training.jl
@@ -0,0 +1,239 @@
+"""
+ createdataloader(data; nuse = 50, device = identity, rng)
+
+Create dataloader that uses a batch of `batchsize` random samples from
+`data` at each evaluation.
+The batch is moved to `device`.
+"""
+create_dataloader_prior(data; batchsize = 50, device = identity, rng) =
+ function dataloader()
+ x, y = data
+ nsample = size(x)[end]
+ d = ndims(x)
+ i = sort(shuffle(rng, 1:nsample)[1:batchsize])
+ xuse = device(Array(selectdim(x, d, i)))
+ yuse = device(Array(selectdim(y, d, i)))
+ xuse, yuse
+ end
+
+"""
+ create_dataloader_post(trajectories; nunroll = 10, device = identity, rng)
+
+Create trajectory dataloader.
+"""
+create_dataloader_post(trajectories; nunroll = 10, device = identity, rng) =
+ function dataloader()
+ (; u, t) = rand(trajectories)
+ nt = length(t)
+ @assert nt ≥ nunroll
+ istart = rand(rng, 1:nt-nunroll)
+ it = istart:istart+nunroll
+ (; u = device.(u[it]), t = t[it])
+ end
+
+"""
+ train(
+ dataloaders,
+ loss,
+ opt,
+ θ;
+ niter = 100,
+ ncallback = 1,
+ callback = (i, θ) -> println("Iteration \$i of \$niter"),
+ )
+
+Update parameters `θ` to minimize `loss(dataloader(), θ)` using the
+optimiser `opt` for `niter` iterations.
+
+Return the a new named tuple `(; opt, θ, callbackstate)` with
+updated state and parameters.
+"""
+function train(
+ dataloaders,
+ loss,
+ opt,
+ θ;
+ niter = 100,
+ ncallback = 1,
+ callback = (state, i, θ) -> println("Iteration $i of $niter"),
+ callbackstate = nothing,
+)
+ for i = 1:niter
+ g = sum(dataloaders) do d
+ b = d()
+ first(gradient(θ -> loss(b, θ), θ))
+ end
+ opt, θ = Optimisers.update(opt, θ, g)
+ if i % ncallback == 0
+ callbackstate = callback(callbackstate, i, θ)
+ end
+ end
+ (; opt, θ, callbackstate)
+end
+
+"""
+ createloss_prior(loss, f)
+
+Wrap loss function `loss(batch, θ)`.
+
+The function `loss` should take inputs like `loss(f, x, y, θ)`.
+"""
+create_loss_prior(loss, f) = ((x, y), θ) -> loss(f, x, y, θ)
+
+"""
+ create_relerr_prior(f, x, y)
+
+Create a-priori error.
+"""
+create_relerr_prior(f, x, y) = θ -> norm(f(x, θ) - y) / norm(y)
+
+"""
+ mean_squared_error(f, x, y, θ; normalize = y -> sum(abs2, y), λ = sqrt(eps(eltype(x))))
+
+Compute MSE between `f(x, θ)` and `y`.
+
+The MSE is further divided by `normalize(y)`.
+"""
+mean_squared_error(f, x, y, θ; normalize = y -> sum(abs2, y), λ = sqrt(eltype(x)(1e-8))) =
+ sum(abs2, f(x, θ) - y) / normalize(y) + λ * sum(abs2, θ)
+
+"""
+ create_loss_post(;
+ setup,
+ method = RKMethods.RK44(; T = eltype(setup.grid.x[1])),
+ psolver,
+ closure,
+ nupdate = 1,
+ projectorder = :last,
+ )
+
+Create a-posteriori loss function.
+"""
+function create_loss_post(;
+ setup,
+ method = RKMethods.RK44(; T = eltype(setup.grid.x[1])),
+ psolver,
+ closure,
+ nupdate = 1,
+ projectorder = :last,
+)
+ closure_model = wrappedclosure(closure, setup)
+ setup = (; setup..., closure_model, projectorder)
+ (; dimension, Iu) = setup.grid
+ D = dimension()
+ function loss_post(data, θ)
+ T = eltype(θ)
+ (; u, t) = data
+ v = u[1]
+ stepper = create_stepper(method; setup, psolver, u = v, t = t[1])
+ loss = zero(eltype(v[1]))
+ for it = 2:length(t)
+ Δt = (t[it] - t[it-1]) / nupdate
+ for isub = 1:nupdate
+ stepper = timestep(method, stepper, Δt; θ)
+ end
+ a, b = T(0), T(0)
+ for α = 1:length(u[1])
+ a += sum(abs2, (stepper.u[α]-u[it][α])[Iu[α]])
+ b += sum(abs2, u[it][α][Iu[α]])
+ end
+ loss += a / b
+ end
+ loss / (length(t) - 1)
+ end
+end
+
+"""
+ create_relerr_post(;
+ data,
+ setup,
+ method = RKMethods.RK44(; T = eltype(setup.grid.x[1])),
+ psolver,
+ closure_model,
+ nupdate = 1,
+ )
+
+Create a-posteriori relative error.
+"""
+function create_relerr_post(;
+ data,
+ setup,
+ method = RKMethods.RK44(; T = eltype(setup.grid.x[1])),
+ psolver,
+ closure_model,
+ nupdate = 1,
+ projectorder = :last,
+)
+ setup = (; setup..., closure_model, projectorder)
+ (; dimension, Iu) = setup.grid
+ D = dimension()
+ (; u, t) = data
+ v = copy.(u[1])
+ cache = ode_method_cache(method, setup, v)
+ function relerr_post(θ)
+ T = eltype(u[1][1])
+ copyto!.(v, u[1])
+ stepper = create_stepper(method; setup, psolver, u = v, t = t[1])
+ e = zero(T)
+ for it = 2:length(t)
+ Δt = (t[it] - t[it-1]) / nupdate
+ for isub = 1:nupdate
+ stepper = timestep!(method, stepper, Δt; θ, cache)
+ end
+ a, b = T(0), T(0)
+ for α = 1:D
+ # a += sum(abs2, (stepper.u[α]-u[it][α])[Iu[α]])
+ # b += sum(abs2, u[it][α][Iu[α]])
+ a += sum(abs2, view(stepper.u[α] - u[it][α], Iu[α]))
+ b += sum(abs2, view(u[it][α], Iu[α]))
+ end
+ e += sqrt(a) / sqrt(b)
+ end
+ e / (length(t) - 1)
+ end
+end
+
+"""
+ create_callback(
+ f,
+ x,
+ y;
+ state = Point2f[],
+ display_each_iteration = false,
+ )
+
+Create convergence plot for relative error between `f(x, θ)` and `y`.
+At each callback, plot is updated and current error is printed.
+
+If `state` is nonempty, it also plots previous convergence.
+
+If not using interactive GLMakie window, set `display_each_iteration` to
+`true`.
+"""
+function create_callback(
+ err;
+ θ,
+ callbackstate = (; θmin = θ, emin = eltype(θ)(Inf), hist = Point2f[]),
+ displayref = true,
+ display_each_iteration = false,
+)
+ istart = isempty(callbackstate.hist) ? 0 : Int(callbackstate.hist[end][1])
+ obs = Observable([Point2f(0, 0)])
+ fig = lines(obs; axis = (; title = "Relative prediction error", xlabel = "Iteration"))
+ displayref && hlines!([1.0f0]; linestyle = :dash)
+ obs[] = callbackstate.hist
+ display(fig)
+ function callback(state, i, θ)
+ e = err(θ)
+ @info "Iteration $i \trelative error: $e"
+ hist = push!(copy(state.hist), Point2f(istart + i, e))
+ obs[] = hist
+ # i < 30 || autolimits!(fig.axis)
+ autolimits!(fig.axis)
+ display_each_iteration && display(fig)
+ state = (; state..., hist)
+ e < state.emin && (state = (; state..., θmin = θ, emin = e))
+ state
+ end
+ (; callbackstate, callback)
+end
diff --git a/lib/NeuralClosure/test/Project.toml b/lib/NeuralClosure/test/Project.toml
new file mode 100644
index 000000000..53c677c2a
--- /dev/null
+++ b/lib/NeuralClosure/test/Project.toml
@@ -0,0 +1,4 @@
+[deps]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+IncompressibleNavierStokes = "5e318141-6589-402b-868d-77d7df8c442e"
+NeuralClosure = "099dac27-d7f2-4047-93d5-0baee36b9c25"
diff --git a/lib/NeuralClosure/test/runtests.jl b/lib/NeuralClosure/test/runtests.jl
new file mode 100644
index 000000000..c4a7c4c87
--- /dev/null
+++ b/lib/NeuralClosure/test/runtests.jl
@@ -0,0 +1,10 @@
+using Aqua
+using IncompressibleNavierStokes
+using NeuralClosure
+
+@testset "NeuralClosure" begin
+ @testset "Aqua" begin
+ @info "Testing code with Aqua"
+ Aqua.test_all(NeuralClosure; ambiguities = false)
+ end
+end
diff --git a/lib/PaperDC/LICENSE b/lib/PaperDC/LICENSE
new file mode 100644
index 000000000..d76a52488
--- /dev/null
+++ b/lib/PaperDC/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Syver Døving Agdestein, Benjamin Sanderse, and contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/lib/PaperDC/Project.toml b/lib/PaperDC/Project.toml
new file mode 100644
index 000000000..16bc477da
--- /dev/null
+++ b/lib/PaperDC/Project.toml
@@ -0,0 +1,24 @@
+name = "PaperDC"
+uuid = "0d06c178-ba13-4c21-b726-6f4489412f0d"
+authors = ["Syver Døving Agdestein "]
+version = "0.1.0"
+
+[deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+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"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+NeuralClosure = "099dac27-d7f2-4047-93d5-0baee36b9c25"
+Observables = "510215fc-4207-5dde-b226-833fc4488ee2"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
diff --git a/lib/PaperDC/README.md b/lib/PaperDC/README.md
new file mode 100644
index 000000000..5b6d4c8c1
--- /dev/null
+++ b/lib/PaperDC/README.md
@@ -0,0 +1,33 @@
+# PaperDC
+
+Scripts for generating results of the paper
+[Discretize first, filter next: learning divergence-consistent closure models for large-eddy simulation](https://arxiv.org/abs/2403.18088).
+
+## Set up environment
+
+From this directory, run:
+
+```sh
+julia --project -e '
+using Pkg
+Pkg.develop([
+ PackageSpec(; path = "../.."),
+ PackageSpec(; path = "../NeuralClosure"),
+])
+Pkg.instantiate()
+'
+```
+
+or interactively from a Julia REPL:
+
+```julia-repl
+julia> ]
+(v1.10) pkg> activate .
+(PaperDC) pkg> dev ../.. ../NeuralClosure
+(PaperDC) pkg> instantiate
+```
+
+Now you can run the scripts in this directory:
+
+- `prioranalysis.jl`: Generate results for section 5.1 "Filtered DNS (2D and 3D)"
+- `postanalysis.jl`: Generate results for section 5.2 "LES (2D)"
diff --git a/lib/PaperDC/postanalysis.jl b/lib/PaperDC/postanalysis.jl
new file mode 100644
index 000000000..2e018d5c4
--- /dev/null
+++ b/lib/PaperDC/postanalysis.jl
@@ -0,0 +1,1143 @@
+# # A-posteriori analysis: Large Eddy Simulation (2D)
+#
+# Generate filtered DNS data, train closure model, compare filters,
+# closure models, and projection orders.
+#
+# The filtered DNS data is saved and can be loaded in a subesequent session.
+# The learned CNN parameters are also saved.
+
+using Adapt
+using GLMakie
+using CairoMakie
+using IncompressibleNavierStokes
+using IncompressibleNavierStokes.RKMethods
+using JLD2
+using LaTeXStrings
+using LinearAlgebra
+using Lux
+using NeuralClosure
+using NNlib
+using Optimisers
+using Random
+using SparseArrays
+using FFTW
+
+# Color palette for consistent theme throughout paper
+palette = (; color = ["#3366cc", "#cc0000", "#669900", "#ff9900"])
+
+# Encode projection order ("close first, then project" etc)
+getorder(i) =
+ if i == 1
+ :first
+ elseif i == 2
+ :last
+ elseif i == 3
+ :second
+ else
+ error("Unknown order: $i")
+ end
+
+# Choose where to put output
+plotdir = "output/postanalysis/plots"
+outdir = "output/postanalysis"
+ispath(plotdir) || mkpath(plotdir)
+ispath(outdir) || mkpath(outdir)
+
+# Random number generator seeds ################################################
+#
+# Use a new RNG with deterministic seed for each code "section"
+# so that e.g. training batch selection does not depend on whether we
+# generated fresh filtered DNS data or loaded existing one (the
+# generation of which would change the state of a global RNG).
+#
+# Note: Using `rng = Random.default_rng()` twice seems to point to the
+# same RNG, and mutating one also mutates the other.
+# `rng = Random.Xoshiro()` creates an independent copy each time.
+#
+# We define all the seeds here so that we don't accidentally type the same seed
+# twice.
+
+seeds = (;
+ dns = 123, # Initial conditions
+ θ₀ = 234, # Initial CNN parameters
+ prior = 345, # A-priori training batch selection
+ post = 456, # A-posteriori training batch selection
+)
+
+# Hardware selection ########################################################
+
+# For running on CPU.
+# Consider reducing the sizes of DNS, LES, and CNN layers if
+# you want to test run on a laptop.
+T = Float32
+ArrayType = Array
+device = identity
+clean() = nothing
+
+# For running on a CUDA compatible GPU
+using LuxCUDA
+using CUDA;
+T = Float32;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+device = x -> adapt(CuArray, x)
+clean() = (GC.gc(); CUDA.reclaim())
+
+# Data generation ###########################################################
+#
+# Create filtered DNS data for training, validation, and testing.
+
+# Random number generator for initial conditions.
+# Important: Created and seeded first, then shared for all initial conditions.
+# After each initial condition generation, it is mutated and creates different
+# IC for the next iteration.
+rng = Random.Xoshiro()
+Random.seed!(rng, seeds.dns)
+
+# Parameters
+get_params(nlesscalar) = (;
+ D = 2,
+ Re = T(10_000),
+ tburn = T(0.05),
+ tsim = T(0.5),
+ Δt = T(5e-5),
+ nles = map(n -> (n, n), nlesscalar), # LES resolutions
+ ndns = (n -> (n, n))(4096), # DNS resolution
+ filters = (FaceAverage(), VolumeAverage()),
+ ArrayType,
+ create_psolver = psolver_spectral,
+ icfunc = (setup, psolver) ->
+ random_field(setup, zero(eltype(setup.grid.x[1])); kp = 20, psolver),
+ rng,
+)
+
+# Get parameters for multiple LES resolutions
+params_train = (; get_params([64, 128, 256])..., tsim = T(0.5), savefreq = 10);
+params_valid = (; get_params([64, 128, 256])..., tsim = T(0.1), savefreq = 40);
+params_test = (; get_params([64, 128, 256, 512, 1024])..., tsim = T(0.1), savefreq = 10);
+
+# Create filtered DNS data
+data_train = [create_les_data(; params_train...) for _ = 1:5];
+data_valid = [create_les_data(; params_valid...) for _ = 1:1];
+data_test = create_les_data(; params_test...);
+
+# Save filtered DNS data
+jldsave("$outdir/data_train.jld2"; data_train)
+jldsave("$outdir/data_valid.jld2"; data_valid)
+jldsave("$outdir/data_test.jld2"; data_test)
+
+# Load filtered DNS data
+data_train = load("$outdir/data_train.jld2", "data_train");
+data_valid = load("$outdir/data_valid.jld2", "data_valid");
+data_test = load("$outdir/data_test.jld2", "data_test");
+
+# Computational time
+data_train[5].comptime
+data_valid[1].comptime
+data_test.comptime
+map(d -> d.comptime, data_train)
+sum(d -> d.comptime, data_train) / 60
+data_test.comptime / 60
+(sum(d -> d.comptime, data_train) + sum(d -> d.comptime, data_valid) + data_test.comptime)
+
+# Build LES setup and assemble operators
+getsetups(params) = [
+ Setup(
+ ntuple(α -> LinRange(T(0), T(1), nles[α] + 1), params.D)...;
+ params.Re,
+ params.ArrayType,
+ ) for nles in params.nles
+]
+setups_train = getsetups(params_train);
+setups_valid = getsetups(params_valid);
+setups_test = getsetups(params_test);
+
+# Example data inspection
+data_train[1].t
+data_train[1].data |> size
+data_train[1].data[1, 1].u[end][1]
+
+# Create input/output arrays for a-priori training (ubar vs c)
+io_train = create_io_arrays(data_train, setups_train);
+io_valid = create_io_arrays(data_valid, setups_valid);
+io_test = create_io_arrays([data_test], setups_test);
+
+# # Save IO arrays
+# jldsave("$outdir/io_train.jld2"; io_train)
+# jldsave("$outdir/io_valid.jld2"; io_valid)
+# jldsave("$outdir/io_test.jld2"; io_test)
+#
+# # Load IO arrays
+# io_train = load("$outdir/io_train.jld2"; "io_train")
+# io_valid = load("$outdir/io_valid.jld2"; "io_valid")
+# io_test = load("$outdir/io_test.jld2"; "io_test")
+
+# Check that data is reasonably bounded
+io_train[1].u |> extrema
+io_train[1].c |> extrema
+io_valid[1].u |> extrema
+io_valid[1].c |> extrema
+io_test[1].u |> extrema
+io_test[1].c |> extrema
+
+# Inspect data (live animation with GLMakie)
+GLMakie.activate!()
+let
+ ig = 2
+ ifil = 1
+ field, setup = data_train[1].data[ig, ifil].u, setups_train[ig]
+ # field, setup = data_valid[1].data[ig, ifil].u, setups_valid[ig];
+ # field, setup = data_test.data[ig, ifil].u, setups_test[ig];
+ u = device.(field[1])
+ o = Observable((; u, t = nothing))
+ # energy_spectrum_plot(o; setup) |> display
+ fieldplot(
+ o;
+ setup,
+ # fieldname = :velocity,
+ # fieldname = 1,
+ ) |> display
+ for i = 1:length(field)
+ o[] = (; o[]..., u = device(field[i]))
+ sleep(0.001)
+ end
+end
+
+# CNN closure model ##########################################################
+
+# Random number generator for initial CNN parameters.
+# All training sessions will start from the same θ₀
+# for a fair comparison.
+rng = Random.Xoshiro()
+Random.seed!(rng, seeds.θ₀)
+
+# # CNN architecture 1
+# mname = "balzac"
+# closure, θ₀ = cnn(;
+# setup = setups_train[1],
+# radii = [2, 2, 2, 2],
+# channels = [20, 20, 20, params_train.D],
+# activations = [leakyrelu, leakyrelu, leakyrelu, identity],
+# use_bias = [true, true, true, false],
+# rng,
+# );
+# closure.chain
+
+# CNN architecture 2
+mname = "rimbaud"
+closure, θ₀ = cnn(;
+ setup = setups_train[1],
+ radii = [2, 2, 2, 2, 2],
+ channels = [24, 24, 24, 24, params_train.D],
+ activations = [tanh, tanh, tanh, tanh, identity],
+ use_bias = [true, true, true, true, false],
+ rng,
+);
+closure.chain
+
+# Save-path for CNN
+savepath = "$outdir/$mname"
+ispath(savepath) || mkpath(savepath)
+
+# Give the CNN a test run
+# Note: Data and parameters are stored on the CPU, and
+# must be moved to the GPU before running (`device`)
+closure(device(io_train[1, 1].u[:, :, :, 1:50]), device(θ₀));
+
+# A-priori training ###########################################################
+#
+# Train one set of CNN parameters for each of the filter types and grid sizes.
+# Save parameters to disk after each run.
+# Plot training progress (for a validation data batch).
+
+priornames = map(CartesianIndices(io_train)) do I
+ ig, ifil = I.I
+ "$savepath/prior_ifilter$(ifil)_igrid$(ig).jld2"
+end
+
+# Train
+let
+ # Random number generator for batch selection
+ rng = Random.Xoshiro()
+ Random.seed!(rng, seeds.prior)
+ ngrid, nfilter = size(io_train)
+ for ifil = 1:nfilter, ig = 1:ngrid
+ clean()
+ starttime = time()
+ println("ig = $ig, ifil = $ifil")
+ d = create_dataloader_prior(io_train[ig, ifil]; batchsize = 50, device, rng)
+ θ = T(1.0e0) * device(θ₀)
+ loss = create_loss_prior(mean_squared_error, closure)
+ opt = Optimisers.setup(Adam(T(1.0e-3)), θ)
+ it = rand(rng, 1:size(io_valid[ig, ifil].u, 4), 50)
+ validset = device(map(v -> v[:, :, :, it], io_valid[ig, ifil]))
+ (; callbackstate, callback) = create_callback(
+ create_relerr_prior(closure, validset...);
+ θ,
+ displayref = true,
+ display_each_iteration = false, # Set to `true` if using CairoMakie
+ )
+ (; opt, θ, callbackstate) = train(
+ [d],
+ loss,
+ opt,
+ θ;
+ niter = 10_000,
+ ncallback = 20,
+ callbackstate,
+ callback,
+ )
+ θ = callbackstate.θmin # Use best θ instead of last θ
+ prior = (; θ = Array(θ), comptime = time() - starttime, callbackstate.hist)
+ jldsave(priorfiles[ig, ifil]; prior)
+ end
+ clean()
+end
+
+# Load learned parameters and training times
+prior = map(f -> load(f)["prior"], priorfiles)
+θ_cnn_prior = [copyto!(device(θ₀), p.θ) for p in prior];
+
+# Check that parameters are within reasonable bounds
+θ_cnn_prior .|> extrema
+
+# Training times
+map(p -> p.comptime, prior)
+map(p -> p.comptime, prior) |> vec
+map(p -> p.comptime, prior) |> sum # Seconds
+map(p -> p.comptime, prior) |> sum |> x -> x / 60 # Minutes
+map(p -> p.comptime, prior) |> sum |> x -> x / 3600 # Hours
+
+# A-posteriori training ######################################################
+#
+# Train one set of CNN parameters for each
+# projection order, filter type and grid size.
+# Save parameters to disk after each combination.
+# Plot training progress (for a validation data batch).
+#
+# The time stepper `RKProject` allows for choosing when to project.
+
+postfiles = map(CartesianIndices((size(io_train)..., 2))) do I
+ ig, ifil, iorder = I.I
+ "$savepath/post_iorder$(iorder)_ifil$(ifil)_ig$(ig).jld2"
+end
+
+# Train
+let
+ # Random number generator for batch selection
+ rng = Random.Xoshiro()
+ Random.seed!(rng, seeds.post)
+ ngrid, nfilter = size(io_train)
+ for iorder = 1:2, ifil = 1:nfilter, ig = 1:ngrid
+ clean()
+ starttime = time()
+ println("iorder = $iorder, ifil = $ifil, ig = $ig")
+ setup = setups_train[ig]
+ psolver = psolver_spectral(setup)
+ loss = IncompressibleNavierStokes.create_loss_post(;
+ setup,
+ psolver,
+ method = RKProject(RK44(; T), getorder(iorder)),
+ closure,
+ nupdate = 2, # Time steps per loss evaluation
+ )
+ data = [(; u = d.data[ig, ifil].u, d.t) for d in data_train]
+ d = create_dataloader_post(data; device, nunroll = 20, rng)
+ θ = copy(θ_cnn_prior[ig, ifil])
+ opt = Optimisers.setup(Adam(T(1.0e-3)), θ)
+ it = 1:30
+ data = data_valid[1]
+ data = (; u = device.(data.data[ig, ifil].u[it]), t = data.t[it])
+ (; callbackstate, callback) = create_callback(
+ create_relerr_post(;
+ data,
+ setup,
+ psolver,
+ method = RKProject(RK44(; T), getorder(iorder)),
+ closure_model = wrappedclosure(closure, setup),
+ nupdate = 2,
+ );
+ θ,
+ displayref = false,
+ )
+ (; opt, θ, callbackstate) =
+ train([d], loss, opt, θ; niter = 2000, ncallback = 10, callbackstate, callback)
+ θ = callbackstate.θmin # Use best θ instead of last θ
+ post = (; θ = Array(θ), comptime = time() - starttime)
+ jldsave(postfiles[iorder, ifil, ig]; post)
+ end
+ clean()
+end
+
+# Load learned parameters and training times
+post = map(f -> load(f)["post"], postfiles);
+θ_cnn_post = [copyto!(device(θ₀), p.θ) for p in post];
+
+# Check that parameters are within reasonable bounds
+θ_cnn_post .|> extrema
+
+# Training times
+map(p -> p.comptime, post)
+map(p -> p.comptime, post) |> x -> reshape(x, 6, 2)
+map(p -> p.comptime, post) ./ 60
+map(p -> p.comptime, post) |> sum
+map(p -> p.comptime, post) |> sum |> x -> x / 60
+map(p -> p.comptime, post) |> sum |> x -> x / 3600
+
+# Train Smagorinsky model ####################################################
+#
+# Use a-posteriori error grid search to determine
+# the optimal Smagorinsky constant.
+# Find one constant for each projection order and filter type. but
+# The constant is shared for all grid sizes, since the filter
+# width (=grid size) is part of the model definition separately.
+
+smag = map(CartesianIndices((size(io_train, 2), 2))) do I
+ starttime = time()
+ ifil, iorder = I.I
+ ngrid = size(io_train, 1)
+ θmin = T(0)
+ emin = T(Inf)
+ isample = 1
+ it = 1:50
+ for θ in LinRange(T(0), T(0.5), 501)
+ e = T(0)
+ for igrid = 1:ngrid
+ println("iorder = $iorder, ifil = $ifil, θ = $θ, igrid = $igrid")
+ projectorder = getorder(iorder)
+ setup = setups_train[igrid]
+ psolver = psolver_spectral(setup)
+ d = data_train[isample]
+ data = (; u = device.(d.data[igrid, ifil].u[it]), t = d.t[it])
+ nupdate = 4
+ err = create_relerr_post(;
+ data,
+ setup,
+ psolver,
+ method = RKProject(RK44(; T), getorder(iorder)),
+ closure_model = smagorinsky_closure(setup),
+ nupdate,
+ )
+ e += err(θ)
+ end
+ e /= ngrid
+ if e < emin
+ emin = e
+ θmin = θ
+ end
+ end
+ (; θ = θmin, comptime = time() - starttime)
+end
+clean()
+
+smag
+
+# Save trained parameters
+jldsave("$outdir/smag.jld2"; smag);
+
+# Load trained parameters
+smag = load("$outdir/smag.jld2")["smag"];
+
+# Extract coefficients
+θ_smag = map(s -> s.θ, smag)
+
+# Computational time
+map(s -> s.comptime, smag)
+map(s -> s.comptime, smag) |> sum
+
+# Compute a-priori errors ###################################################
+#
+# Note that it is still interesting to compute the a-priori errors for the
+# a-posteriori trained CNN.
+
+eprior = let
+ prior = zeros(T, 3, 2)
+ post = zeros(T, 3, 2, 2)
+ for ig = 1:3, ifil = 1:2
+ println("ig = $ig, ifil = $ifil")
+ testset = device(io_test[ig, ifil])
+ err = create_relerr_prior(closure, testset...)
+ prior[ig, ifil] = err(θ_cnn_prior[ig, ifil])
+ for iorder = 1:2
+ post[ig, ifil, iorder] = err(θ_cnn_post[ig, ifil, iorder])
+ end
+ end
+ (; prior, post)
+end
+clean()
+
+eprior.prior
+eprior.post
+
+eprior.prior |> x -> reshape(x, :) |> x -> round.(x; digits = 2)
+eprior.post |> x -> reshape(x, :, 2) |> x -> round.(x; digits = 2)
+
+# Compute a-posteriori errors #################################################
+
+(; e_nm, e_smag, e_cnn, e_cnn_post) = let
+ e_nm = zeros(T, size(data_test.data)...)
+ e_smag = zeros(T, size(data_test.data)..., 2)
+ e_cnn = zeros(T, size(data_test.data)..., 2)
+ e_cnn_post = zeros(T, size(data_test.data)..., 2)
+ for iorder = 1:2, ifil = 1:2, ig = 1:size(data_test.data, 1)
+ println("iorder = $iorder, ifil = $ifil, ig = $ig")
+ projectorder = getorder(iorder)
+ setup = setups_test[ig]
+ psolver = psolver_spectral(setup)
+ data = (; u = device.(data_test.data[ig, ifil].u), t = data_test.t)
+ nupdate = 2
+ # No model
+ # Only for closurefirst, since projectfirst is the same
+ if iorder == 2
+ err = create_relerr_post(; data, setup, psolver, closure_model = nothing, nupdate)
+ e_nm[ig, ifil] = err(nothing)
+ end
+ # Smagorinsky
+ err = create_relerr_post(;
+ data,
+ setup,
+ psolver,
+ method = RKProject(RK44(; T), getorder(iorder)),
+ closure_model = smagorinsky_closure(setup),
+ nupdate,
+ )
+ e_smag[ig, ifil, iorder] = err(θ_smag[ifil, iorder])
+ # CNN
+ # Only the first grids are trained for
+ if ig ≤ size(data_train[1].data, 1)
+ err = create_relerr_post(;
+ data,
+ setup,
+ psolver,
+ method = RKProject(RK44(; T), getorder(iorder)),
+ closure_model = wrappedclosure(closure, setup),
+ nupdate,
+ )
+ e_cnn[ig, ifil, iorder] = err(θ_cnn_prior[ig, ifil])
+ e_cnn_post[ig, ifil, iorder] = err(θ_cnn_post[ig, ifil, iorder])
+ end
+ end
+ (; e_nm, e_smag, e_cnn, e_cnn_post)
+end
+clean()
+
+round.(
+ [e_nm[:] reshape(e_smag, :, 2) reshape(e_cnn, :, 2) reshape(e_cnn_post, :, 2)][
+ [1:3; 6:8],
+ :,
+ ];
+ sigdigits = 2,
+)
+
+# Plot a-priori errors ########################################################
+
+# Better for PDF export
+CairoMakie.activate!()
+
+fig = with_theme(; palette) do
+ nles = [n[1] for n in params_test.nles][1:3]
+ ifil = 1
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xscale = log10,
+ xticks = nles,
+ xlabel = "Resolution",
+ title = "Relative a-priori error $(ifil == 1 ? " (FA)" : " (VA)")",
+ )
+ linestyle = :solid
+ label = "No closure"
+ scatterlines!(
+ nles,
+ ones(T, length(nles));
+ color = Cycled(1),
+ linestyle,
+ marker = :circle,
+ label,
+ )
+ label = "CNN (Lprior)"
+ scatterlines!(
+ nles,
+ eprior.prior[:, ifil];
+ color = Cycled(2),
+ linestyle,
+ marker = :utriangle,
+ label,
+ )
+ label = "CNN (Lpost, DIF)"
+ scatterlines!(
+ nles,
+ eprior.post[:, ifil, 1];
+ color = Cycled(3),
+ linestyle,
+ marker = :rect,
+ label,
+ )
+ label = "CNN (Lpost, DCF)"
+ scatterlines!(
+ nles,
+ eprior.post[:, ifil, 2];
+ color = Cycled(4),
+ linestyle,
+ marker = :diamond,
+ label,
+ )
+ axislegend(; position = :lb)
+ ylims!(ax, (T(-0.05), T(1.05)))
+ name = "$plotdir/convergence"
+ ispath(name) || mkpath(name)
+ save("$name/$(mname)_prior_ifilter$ifil.pdf", fig)
+ fig
+end
+
+# Plot a-posteriori errors ###################################################
+
+# Better for PDF export
+CairoMakie.activate!()
+
+with_theme(; palette) do
+ iorder = 2
+ lesmodel = iorder == 1 ? "DIF" : "DCF"
+ ntrain = size(data_train[1].data, 1)
+ nles = [n[1] for n in params_test.nles][1:ntrain]
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xscale = log10,
+ yscale = log10,
+ xticks = nles,
+ xlabel = "Resolution",
+ title = "Relative error ($lesmodel)",
+ )
+ for ifil = 1:2
+ linestyle = ifil == 1 ? :solid : :dash
+ label = "No closure"
+ ifil == 2 && (label = nothing)
+ scatterlines!(
+ nles,
+ e_nm[1:ntrain, ifil];
+ color = Cycled(1),
+ linestyle,
+ marker = :circle,
+ label,
+ )
+ end
+ for ifil = 1:2
+ linestyle = ifil == 1 ? :solid : :dash
+ label = "Smagorinsky"
+ ifil == 2 && (label = nothing)
+ scatterlines!(
+ nles,
+ e_smag[1:ntrain, ifil, iorder];
+ color = Cycled(2),
+ linestyle,
+ marker = :utriangle,
+ label,
+ )
+ end
+ for ifil = 1:2
+ linestyle = ifil == 1 ? :solid : :dash
+ label = "CNN (prior)"
+ ifil == 2 && (label = nothing)
+ scatterlines!(
+ nles[1:ntrain],
+ e_cnn[1:ntrain, ifil, iorder];
+ color = Cycled(3),
+ linestyle,
+ marker = :rect,
+ label,
+ )
+ end
+ for ifil = 1:2
+ linestyle = ifil == 1 ? :solid : :dash
+ label = "CNN (post)"
+ ifil == 2 && (label = nothing)
+ scatterlines!(
+ nles[1:ntrain],
+ e_cnn_post[1:ntrain, ifil, iorder];
+ color = Cycled(4),
+ linestyle,
+ marker = :diamond,
+ label,
+ )
+ end
+ axislegend(; position = :lb)
+ ylims!(ax, (T(0.025), T(1.00)))
+ name = "$plotdir/convergence"
+ ispath(name) || mkpath(name)
+ save("$name/$(mname)_iorder$iorder.pdf", fig)
+ fig
+end
+
+# Energy evolution ###########################################################
+#
+# Compute total kinetic energy as a function of time.
+
+kineticenergy = let
+ clean()
+ ngrid, nfilter = size(io_train)
+ ke_ref = fill(zeros(T, 0), ngrid, nfilter)
+ ke_nomodel = fill(zeros(T, 0), ngrid, nfilter)
+ ke_smag = fill(zeros(T, 0), ngrid, nfilter, 2)
+ ke_cnn_prior = fill(zeros(T, 0), ngrid, nfilter, 2)
+ ke_cnn_post = fill(zeros(T, 0), ngrid, nfilter, 2)
+ for iorder = 1:2, ifil = 1:nfilter, ig = 1:ngrid
+ println("iorder = $iorder, ifil = $ifil, ig = $ig")
+ setup = setups_test[ig]
+ psolver = psolver_spectral(setup)
+ t = data_test.t
+ ustart = data_test.data[ig, ifil].u[1] |> device
+ tlims = (t[1], t[end])
+ nupdate = 2
+ Δt = (t[2] - t[1]) / nupdate
+ T = eltype(ustart[1])
+ ewriter = processor() do state
+ ehist = zeros(T, 0)
+ on(state) do (; u, n)
+ if n % nupdate == 0
+ e = IncompressibleNavierStokes.total_kinetic_energy(u, setup)
+ push!(ehist, e)
+ end
+ end
+ state[] = state[] # Compute initial energy
+ ehist
+ end
+ processors = (; ewriter)
+ if iorder == 1
+ # Does not depend on projection order
+ ke_ref[ig, ifil] = map(
+ u -> IncompressibleNavierStokes.total_kinetic_energy(device(u), setup),
+ data_test.data[ig, ifil].u,
+ )
+ ke_nomodel[ig, ifil] =
+ solve_unsteady(; setup, ustart, tlims, Δt, processors, psolver)[2].ewriter
+ end
+ ke_smag[ig, ifil, iorder] =
+ solve_unsteady(;
+ (;
+ setup...,
+ projectorder = getorder(iorder),
+ closure_model = smagorinsky_closure(setup),
+ ),
+ ustart,
+ tlims,
+ Δt,
+ processors,
+ psolver,
+ θ = θ_smag[ifil, iorder],
+ )[2].ewriter
+ ke_cnn_prior[ig, ifil, iorder] =
+ solve_unsteady(;
+ (;
+ setup...,
+ projectorder = getorder(iorder),
+ closure_model = wrappedclosure(closure, setup),
+ ),
+ ustart,
+ tlims,
+ Δt,
+ processors,
+ psolver,
+ θ = θ_cnn_prior[ig, ifil],
+ )[2].ewriter
+ ke_cnn_post[ig, ifil, iorder] =
+ solve_unsteady(;
+ (;
+ setup...,
+ projectorder = getorder(iorder),
+ closure_model = wrappedclosure(closure, setup),
+ ),
+ ustart,
+ tlims,
+ Δt,
+ processors,
+ psolver,
+ θ = θ_cnn_post[ig, ifil, iorder],
+ )[2].ewriter
+ end
+ (; ke_ref, ke_nomodel, ke_smag, ke_cnn_prior, ke_cnn_post)
+end;
+clean();
+
+# Plot energy evolution ########################################################
+
+# Better for PDF export
+CairoMakie.activate!()
+
+with_theme(; palette) do
+ t = data_test.t
+ for iorder = 1:2, ifil = 1:2, igrid = 1:3
+ println("iorder = $iorder, ifil = $ifil, igrid = $igrid")
+ lesmodel = iorder == 1 ? "DIF" : "DCF"
+ fil = ifil == 1 ? "FA" : "VA"
+ nles = params_test.nles[igrid]
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xlabel = "t",
+ ylabel = "E(t)",
+ title = "Kinetic energy: $lesmodel, $fil",
+ )
+ lines!(
+ ax,
+ t,
+ kineticenergy.ke_ref[igrid, ifil];
+ color = Cycled(1),
+ linestyle = :dash,
+ label = "Reference",
+ )
+ lines!(
+ ax,
+ t,
+ kineticenergy.ke_nomodel[igrid, ifil];
+ color = Cycled(1),
+ label = "No closure",
+ )
+ lines!(
+ ax,
+ t,
+ kineticenergy.ke_smag[igrid, ifil, iorder];
+ color = Cycled(2),
+ label = "Smagorinsky",
+ )
+ lines!(
+ ax,
+ t,
+ kineticenergy.ke_cnn_prior[igrid, ifil, iorder];
+ color = Cycled(3),
+ label = "CNN (prior)",
+ )
+ lines!(
+ ax,
+ t,
+ kineticenergy.ke_cnn_post[igrid, ifil, iorder];
+ color = Cycled(4),
+ label = "CNN (post)",
+ )
+ iorder == 1 && axislegend(; position = :lt)
+ iorder == 2 && axislegend(; position = :lb)
+ name = "$plotdir/energy_evolution/$mname/"
+ ispath(name) || mkpath(name)
+ save("$(name)/iorder$(iorder)_ifilter$(ifil)_igrid$(igrid).pdf", fig)
+ end
+end
+
+# Compute Divergence ##########################################################
+#
+# Compute divergence as a function of time.
+
+divs = let
+ clean()
+ ngrid, nfilter = size(io_train)
+ d_ref = fill(zeros(T, 0), ngrid, nfilter)
+ d_nomodel = fill(zeros(T, 0), ngrid, nfilter, 3)
+ d_smag = fill(zeros(T, 0), ngrid, nfilter, 3)
+ d_cnn_prior = fill(zeros(T, 0), ngrid, nfilter, 3)
+ d_cnn_post = fill(zeros(T, 0), ngrid, nfilter, 3)
+ for iorder = 1:3, ifil = 1:nfilter, ig = 1:ngrid
+ println("iorder = $iorder, ifil = $ifil, ig = $ig")
+ setup = setups_test[ig]
+ psolver = psolver_spectral(setup)
+ t = data_test.t
+ ustart = data_test.data[ig, ifil].u[1] |> device
+ tlims = (t[1], t[end])
+ nupdate = 2
+ Δt = (t[2] - t[1]) / nupdate
+ T = eltype(ustart[1])
+ dwriter = processor() do state
+ div = fill!(similar(setup.grid.x[1], setup.grid.N), 0)
+ dhist = zeros(T, 0)
+ on(state) do (; u, n)
+ if n % nupdate == 0
+ IncompressibleNavierStokes.divergence!(div, u, setup)
+ d = view(div, setup.grid.Ip)
+ d = sum(abs2, d) / length(d)
+ d = sqrt(d)
+ push!(dhist, d)
+ end
+ end
+ state[] = state[] # Compute initial divergence
+ dhist
+ end
+ if iorder == 1
+ # Does not depend on projection order
+ d_ref[ig, ifil] = map(data_test.data[ig, ifil].u) do u
+ u = device(u)
+ div = IncompressibleNavierStokes.divergence(u, setup)
+ d = view(div, setup.grid.Ip)
+ d = sum(abs2, d) / length(d)
+ d = sqrt(d)
+ end
+ end
+ s(closure_model, θ) =
+ solve_unsteady(;
+ (; setup..., closure_model),
+ ustart,
+ tlims,
+ method = RKProject(RK44(; T), getorder(iorder)),
+ Δt,
+ processors = (; dwriter),
+ psolver,
+ θ,
+ )[2].dwriter
+ iorder_use = iorder == 3 ? 2 : iorder
+ d_nomodel[ig, ifil, iorder] = s(nothing, nothing)
+ d_smag[ig, ifil, iorder] =
+ s(smagorinsky_closure(setup), θ_smag[ifil, iorder_use])
+ d_cnn_prior[ig, ifil, iorder] =
+ s(wrappedclosure(closure, setup), θ_cnn_prior[ig, ifil])
+ d_cnn_post[ig, ifil, iorder] =
+ s(wrappedclosure(closure, setup), θ_cnn_post[ig, ifil, iorder_use])
+ end
+ (; d_ref, d_nomodel, d_smag, d_cnn_prior, d_cnn_post)
+end;
+clean();
+
+# Check that divergence is within reasonable bounds
+divs.d_ref .|> extrema
+divs.d_nomodel .|> extrema
+divs.d_smag .|> extrema
+divs.d_cnn_prior .|> extrema
+divs.d_cnn_post .|> extrema
+
+# Plot Divergence #############################################################
+
+# Better for PDF export
+CairoMakie.activate!()
+
+with_theme(;
+ # fontsize = 20,
+ palette,
+) do
+ t = data_test.t
+ # for islog in (true, false)
+ for islog in (false,)
+ for iorder = 1:2, ifil = 1:2, igrid = 1:3
+ println("iorder = $iorder, ifil = $ifil, igrid = $igrid")
+ lesmodel = if iorder == 1
+ "DIF"
+ elseif iorder == 2
+ "DCF"
+ elseif iorder == 3
+ "DCF-RHS"
+ end
+ fil = ifil == 1 ? "FA" : "VA"
+ nles = params_test.nles[igrid]
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ yscale = islog ? log10 : identity,
+ xlabel = "t",
+ title = "Divergence: $lesmodel, $fil, $nles",
+ )
+ lines!(ax, t, divs.d_nomodel[igrid, ifil, iorder]; label = "No closure")
+ lines!(ax, t, divs.d_smag[igrid, ifil, iorder]; label = "Smagorinsky")
+ lines!(ax, t, divs.d_cnn_prior[igrid, ifil, iorder]; label = "CNN (prior)")
+ lines!(ax, t, divs.d_cnn_post[igrid, ifil, iorder]; label = "CNN (post)")
+ lines!(
+ ax,
+ t,
+ divs.d_ref[igrid, ifil];
+ color = Cycled(1),
+ linestyle = :dash,
+ label = "Reference",
+ )
+ iorder == 2 && ifil == 1 && axislegend(; position = :rt)
+ islog && ylims!(ax, (T(1e-6), T(1e3)))
+ name = "$plotdir/divergence/$mname/$(islog ? "log" : "lin")"
+ ispath(name) || mkpath(name)
+ save("$(name)/iorder$(iorder)_ifilter$(ifil)_igrid$(igrid).pdf", fig)
+ end
+ end
+end
+
+# Solutions at final time ####################################################
+
+ufinal = let
+ ngrid, nfilter = size(io_train)
+ temp = ntuple(α -> zeros(T, 0, 0), 2)
+ u_ref = fill(temp, ngrid, nfilter)
+ u_nomodel = fill(temp, ngrid, nfilter)
+ u_smag = fill(temp, ngrid, nfilter, 2)
+ u_cnn_prior = fill(temp, ngrid, nfilter, 2)
+ u_cnn_post = fill(temp, ngrid, nfilter, 2)
+ for iorder = 1:2, ifil = 1:nfilter, igrid = 1:ngrid
+ clean()
+ println("iorder = $iorder, ifil = $ifil, igrid = $igrid")
+ t = data_test.t
+ setup = setups_test[igrid]
+ psolver = psolver_spectral(setup)
+ ustart = data_test.data[igrid, ifil].u[1] |> device
+ tlims = (t[1], t[end])
+ nupdate = 2
+ Δt = (t[2] - t[1]) / nupdate
+ T = eltype(ustart[1])
+ s(closure_model, θ) =
+ solve_unsteady(;
+ (; setup..., closure_model),
+ ustart,
+ tlims,
+ method = RKProject(RK44(; T), getorder(iorder)),
+ Δt,
+ psolver,
+ θ,
+ )[1].u .|> Array
+ if iorder == 1
+ # Does not depend on projection order
+ u_ref[igrid, ifil] = data_test.data[igrid, ifil].u[end]
+ u_nomodel[igrid, ifil] = s(nothing, nothing)
+ end
+ u_smag[igrid, ifil, iorder] =
+ s(smagorinsky_closure(setup), θ_smag[ifil, iorder])
+ u_cnn_prior[igrid, ifil, iorder] =
+ s(wrappedclosure(closure, setup), θ_cnn_prior[igrid, ifil])
+ u_cnn_post[igrid, ifil, iorder] =
+ s(wrappedclosure(closure, setup), θ_cnn_post[igrid, ifil, iorder])
+ end
+ (; u_ref, u_nomodel, u_smag, u_cnn_prior, u_cnn_post)
+end;
+clean();
+
+# # Save solution
+# jldsave("$savepath/ufinal.jld2"; ufinal)
+#
+# # Load solution
+# ufinal = load("$savepath/ufinal.jld2")["ufinal"];
+
+# Plot spectra ###############################################################
+#
+# Plot kinetic energy spectra at final time.
+
+# Better for PDF export
+CairoMakie.activate!()
+
+fig = with_theme(; palette) do
+ for iorder = 1:2, ifil = 1:2, igrid = 1:3
+ println("iorder = $iorder, ifil = $ifil, igrid = $igrid")
+ lesmodel = iorder == 1 ? "DIF" : "DCF"
+ fil = ifil == 1 ? "FA" : "VA"
+ nles = params_test.nles[igrid]
+ setup = setups_test[igrid]
+ fields =
+ [
+ ufinal.u_ref[igrid, ifil],
+ ufinal.u_nomodel[igrid, ifil],
+ ufinal.u_smag[igrid, ifil, iorder],
+ ufinal.u_cnn_prior[igrid, ifil, iorder],
+ ufinal.u_cnn_post[igrid, ifil, iorder],
+ ] .|> device
+ (; Ip) = setup.grid
+ (; A, κ, K) = IncompressibleNavierStokes.spectral_stuff(setup)
+ specs = map(fields) do u
+ up = u
+ e = sum(up) do u
+ u = u[Ip]
+ uhat = fft(u)[ntuple(α -> 1:K[α], 2)...]
+ abs2.(uhat) ./ (2 * prod(size(u))^2)
+ end
+ e = A * reshape(e, :)
+ # e = max.(e, eps(T)) # Avoid log(0)
+ ehat = Array(e)
+ end
+ kmax = maximum(κ)
+ # Build inertial slope above energy
+ krange = [T(16), T(κ[end])]
+ slope, slopelabel = -T(3), L"$\kappa^{-3}"
+ slopeconst = maximum(specs[1] ./ κ .^ slope)
+ offset = 3
+ inertia = offset .* slopeconst .* krange .^ slope
+ # Nice ticks
+ logmax = round(Int, log2(kmax + 1))
+ xticks = T(2) .^ (0:logmax)
+ # Make plot
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xticks,
+ xlabel = "κ",
+ xscale = log10,
+ yscale = log10,
+ limits = (1, kmax, T(1e-8), T(1)),
+ title = "Kinetic energy: $lesmodel, $fil",
+ )
+ lines!(ax, κ, specs[2]; color = Cycled(1), label = "No model")
+ lines!(ax, κ, specs[3]; color = Cycled(2), label = "Smagorinsky")
+ lines!(ax, κ, specs[4]; color = Cycled(3), label = "CNN (prior)")
+ lines!(ax, κ, specs[5]; color = Cycled(4), label = "CNN (post)")
+ lines!(ax, κ, specs[1]; color = Cycled(1), linestyle = :dash, label = "Reference")
+ lines!(ax, krange, inertia; color = Cycled(1), label = slopelabel, linestyle = :dot)
+ axislegend(ax; position = :cb)
+ autolimits!(ax)
+ ylims!(ax, (T(1e-3), T(0.35)))
+ name = "$plotdir/energy_spectra/$mname"
+ ispath(name) || mkpath(name)
+ save("$(name)/iorder$(iorder)_ifilter$(ifil)_igrid$(igrid).pdf", fig)
+ end
+end
+clean();
+
+# Plot fields ################################################################
+
+# Export to PNG, otherwise each volume gets represented
+# as a separate rectangle in the PDF
+# (takes time to load in the article PDF)
+GLMakie.activate!()
+
+with_theme(; fontsize = 25, palette) do
+ # Reference box for eddy comparison
+ x1 = 0.3
+ x2 = 0.5
+ y1 = 0.5
+ y2 = 0.7
+ box = [
+ Point2f(x1, y1),
+ Point2f(x2, y1),
+ Point2f(x2, y2),
+ Point2f(x1, y2),
+ Point2f(x1, y1),
+ ]
+ path = "$plotdir/les_fields/$mname"
+ ispath(path) || mkpath(path)
+ for iorder = 1:2, ifil = 1:2, igrid = 1:3
+ setup = setups_test[igrid]
+ name = "$path/iorder$(iorder)_ifilter$(ifil)_igrid$(igrid)"
+ lesmodel = iorder == 1 ? "DIF" : "DCF"
+ fil = ifil == 1 ? "FA" : "VA"
+ nles = params_test.nles[igrid]
+ function makeplot(u, title, suffix)
+ fig = fieldplot(
+ (; u, t = T(0));
+ setup,
+ title,
+ docolorbar = false,
+ size = (500, 500),
+ )
+ lines!(box; linewidth = 5, color = Cycled(2)) # Red in palette
+ fname = "$(name)_$(suffix).png"
+ save(fname, fig)
+ # run(`convert $fname -trim $fname`) # Requires imagemagick
+ end
+ iorder == 2 &&
+ makeplot(device(ufinal.u_ref[igrid, ifil]), "Reference, $fil, $nles", "ref")
+ iorder == 2 && makeplot(
+ device(ufinal.u_nomodel[igrid, ifil]),
+ "No closure, $fil, $nles",
+ "nomodel",
+ )
+ makeplot(
+ device(ufinal.u_smag[igrid, ifil, iorder]),
+ "Smagorinsky, $lesmodel, $fil, $nles",
+ "smag",
+ )
+ makeplot(
+ device(ufinal.u_cnn_prior[igrid, ifil, iorder]),
+ "CNN (prior), $lesmodel, $fil, $nles",
+ "prior",
+ )
+ makeplot(
+ device(ufinal.u_cnn_post[igrid, ifil, iorder]),
+ "CNN (post), $lesmodel, $fil, $nles",
+ "post",
+ )
+ end
+end
diff --git a/lib/PaperDC/prioranalysis.jl b/lib/PaperDC/prioranalysis.jl
new file mode 100644
index 000000000..16998f5f4
--- /dev/null
+++ b/lib/PaperDC/prioranalysis.jl
@@ -0,0 +1,289 @@
+# # A-priori analysis: Filtered DNS (2D or 3D)
+#
+# Generate filtered DNS data and inspect its properties.
+
+using CairoMakie
+using FFTW
+using GLMakie
+using IncompressibleNavierStokes
+using IncompressibleNavierStokes: momentum, project, apply_bc_u!, spectral_stuff
+using NeuralClosure
+using PaperDC
+using Printf
+using Random
+
+# Put output in PaperDC/output
+cd(@__DIR__)
+output = "output"
+
+# For running on CPU
+ArrayType = Array
+clean() = nothing
+
+# For running on GPU
+using CUDA;
+CUDA.allowscalar(false);
+ArrayType = CuArray;
+clean() = (GC.gc(); CUDA.reclaim()) # This seems to be needed to free up memory
+
+# 2D setup
+D = 2
+T = Float64;
+ndns = 4096
+Re = T(10_000)
+kp = 20
+Δt = T(5e-5)
+filterdefs = [
+ (FaceAverage(), 64),
+ (FaceAverage(), 128),
+ (FaceAverage(), 256),
+ (VolumeAverage(), 64),
+ (VolumeAverage(), 128),
+ (VolumeAverage(), 256),
+]
+
+# 3D setup
+D = 3
+T, ndns = Float64, 512 # Works on a 40GB A100 GPU. Use Float32 and smaller n for less memory.
+Re = T(2_000)
+kp = 5
+Δt = T(1e-4)
+filterdefs = [
+ (FaceAverage(), 32),
+ (FaceAverage(), 64),
+ (FaceAverage(), 128),
+ (VolumeAverage(), 32),
+ (VolumeAverage(), 64),
+ (VolumeAverage(), 128),
+]
+
+# Setup
+lims = T(0), T(1)
+dns = let
+ setup = Setup(ntuple(α -> LinRange(lims..., ndns + 1), D)...; Re, ArrayType)
+ psolver = psolver_spectral(setup)
+ (; setup, psolver)
+end;
+filters = map(filterdefs) do (Φ, nles)
+ compression = ndns ÷ nles
+ setup = Setup(ntuple(α -> LinRange(lims..., nles + 1), D)...; Re, ArrayType)
+ psolver = psolver_spectral(setup)
+ (; setup, Φ, compression, psolver)
+end;
+
+# Create random initial conditions
+rng = Random.seed!(Random.default_rng(), 12345)
+ustart = random_field(dns.setup, T(0); kp, dns.psolver, rng);
+clean()
+
+# Solve unsteady problem
+@time state, outputs = solve_unsteady(;
+ dns.setup,
+ ustart,
+ tlims = (T(0), T(1e-1)),
+ Δt,
+ docopy = true, # leave initial conditions unchanged, false to free up memory
+ dns.psolver,
+ processors = (
+ obs = observe_u(dns.setup, dns.psolver, filters; nupdate = 20),
+ log = timelogger(; nupdate = 5),
+ ),
+);
+clean()
+
+# Plot 2D fields #####################################################
+
+D == 2 && with_theme(; fontsize = 25) do
+ # Compute quantities
+ fil = filters[2]
+ apply_bc_u!(state.u, T(0), dns.setup)
+ v = fil.Φ(state.u, fil.setup, fil.compression)
+ apply_bc_u!(v, T(0), fil.setup)
+ Fv = momentum(v, nothing, T(0), fil.setup)
+ apply_bc_u!(Fv, T(0), fil.setup)
+ PFv = project(Fv, fil.setup; psolver = fil.psolver)
+ apply_bc_u!(PFv, T(0), fil.setup)
+ F = momentum(state.u, nothing, T(0), dns.setup)
+ apply_bc_u!(F, T(0), dns.setup)
+ PF = project(F, dns.setup; dns.psolver)
+ apply_bc_u!(PF, T(0), dns.setup)
+ ΦPF = fil.Φ(PF, fil.setup, fil.compression)
+ apply_bc_u!(ΦPF, T(0), fil.setup)
+ c = ΦPF .- PFv
+ apply_bc_u!(c, T(0), fil.setup)
+
+ # Make plots
+ path = "$output/priorfields"
+ ispath(path) || mkpath(path)
+ makeplot(field, setup, title, name) = save(
+ "$path/$name.png",
+ fieldplot(
+ (; u = field, t = T(0));
+ setup,
+ title,
+ docolorbar = false,
+ size = (500, 500),
+ ),
+ )
+ makeplot(u₀, dns.setup, "u₀", "ustart")
+ makeplot(state.u, dns.setup, "u", "u")
+ makeplot(v, fil.setup, "ū", "v")
+ makeplot(PF, dns.setup, "PF(u)", "PFu")
+ makeplot(PFv, fil.setup, "P̄F̄(ū)", "PFv")
+ makeplot(ΦPF, fil.setup, "ΦPF(u)", "PhiPFu")
+ makeplot(c, fil.setup, "c(u, ū)", "c")
+end
+
+# Plot 3D fields #####################################################
+
+# Contour plots in 3D only work with GLMakie.
+# For using GLMakie on headless servers, see
+#
+GLMakie.activate!()
+
+# Make plots
+D == 3 && with_theme() do
+ path = "$output/priorfields/lambda2"
+ ispath(path) || mkpath(path)
+ function makeplot(field, setup, name)
+ name = "$path/$name.png"
+ save(
+ name,
+ fieldplot(
+ (; u = field, t = T(0));
+ setup,
+ fieldname = :eig2field,
+ levels = LinRange(T(4), T(12), 10),
+ docolorbar = false,
+ size = (600, 600),
+ ),
+ )
+ try
+ # Trim whitespace with ImageMagick
+ run(`convert $name -trim $name`)
+ catch e
+ @warn """
+ ImageMagick not found.
+ Skipping image trimming.
+ Install from .
+ """
+ end
+ end
+ makeplot(u₀, dns.setup, "Re$(Int(Re))_start") # Requires docopy = true in solve
+ makeplot(state.u, dns.setup, "Re$(Int(Re))_end")
+ i = 3
+ makeplot(
+ filters[i].Φ(state.u, filters[i].setup, filters[i].compression),
+ filters[i].setup,
+ "Re$(Int(Re))_end_filtered",
+ )
+end
+
+# Compute average quantities #########################################
+
+let
+ path = "$output/prioranalysis"
+ ispath(path) || mkpath(path)
+ open("$path/averages_$(D)D.txt", "w") do io
+ println(io, "Φ\t\tM\tDu\tPv\tPc\tc")
+ for o in outputs.obs
+ nt = length(o.t)
+ Dv = sum(o.Dv) / nt
+ Pc = sum(o.Pc) / nt
+ Pv = sum(o.Pv) / nt
+ c = sum(o.c) / nt
+ @printf(
+ io,
+ "%s\t%d^%d\t%.2g\t%.2g\t%.2g\t%.2g\n",
+ # "%s &\t\$%d^%d\$ &\t\$%.2g\$ &\t\$%.2g\$ &\t\$%.2g\$ &\t\$%.2g\$\n",
+ typeof(o.Φ),
+ o.Mα,
+ D,
+ Dv,
+ Pv,
+ Pc,
+ c
+ )
+ end
+ end
+end
+
+# Plot spectra #######################################################
+
+# To free up memory in 3D (remove psolver_spectral FFT arrays)
+dns = (; dns.setup)
+filters = map(filters) do f
+ (; f.Φ, f.setup, f.compression)
+end
+fig = lines([1, 2, 3])
+clean()
+
+# Plot predicted spectra
+CairoMakie.activate!()
+with_theme(; palette = (; color = ["#3366cc", "#cc0000", "#669900", "#ffcc00"])) do
+ fields = [state.u, u₀, (f.Φ(state.u, f.setup, f.compression) for f in filters)...]
+ setups = [dns.setup, dns.setup, (f.setup for f in filters)...]
+ specs = map(fields, setups) do u, setup
+ clean() # Free up memory
+ (; dimension, xp, Ip) = setup.grid
+ T = eltype(xp[1])
+ D = dimension()
+ K = size(Ip) .÷ 2
+ up = u
+ e = sum(up) do u
+ u = u[Ip]
+ uhat = fft(u)[ntuple(α -> 1:K[α], D)...]
+ abs2.(uhat) ./ (2 * prod(size(u))^2)
+ end
+ (; A, κ, K) = spectral_stuff(setup) # Requires some memory
+ e = A * reshape(e, :) # Dyadic binning
+ ehat = Array(e) # Store spectrum on CPU
+ (; κ, ehat)
+ end
+ kmax = maximum(specs[1].κ)
+ # Build inertial slope above energy
+ krange, slope, slopelabel = if D == 2
+ [T(16), T(128)], -T(3), L"$\kappa^{-3}"
+ elseif D == 3
+ [T(16), T(100)], -T(5 / 3), L"$\kappa^{-5/3}"
+ end
+ slopeconst = maximum(specs[1].ehat ./ specs[1].κ .^ slope)
+ offset = D == 2 ? 3 : 2
+ inertia = offset .* slopeconst .* krange .^ slope
+ # Nice ticks
+ logmax = round(Int, log2(kmax + 1))
+ xticks = T(2) .^ (0:logmax)
+ # Make plot
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xticks,
+ xlabel = "κ",
+ xscale = log10,
+ yscale = log10,
+ limits = (1, kmax, T(1e-8), T(1)),
+ title = "Kinetic energy ($(D)D)",
+ )
+ plotparts(i) = specs[i].κ, specs[i].ehat
+ lines!(ax, plotparts(1)...; color = Cycled(1), label = "DNS")
+ lines!(ax, plotparts(2)...; color = Cycled(4), label = "DNS, t = 0")
+ lines!(ax, plotparts(3)...; color = Cycled(2), label = "Filtered DNS (FA)")
+ lines!(ax, plotparts(4)...; color = Cycled(2))
+ lines!(ax, plotparts(5)...; color = Cycled(2))
+ lines!(ax, plotparts(6)...; color = Cycled(3), label = "Filtered DNS (VA)")
+ lines!(ax, plotparts(7)...; color = Cycled(3))
+ lines!(ax, plotparts(8)...; color = Cycled(3))
+ lines!(ax, krange, inertia; color = Cycled(1), label = slopelabel, linestyle = :dash)
+ axislegend(ax; position = :lb)
+ autolimits!(ax)
+ if D == 2
+ limits!(ax, (T(0.8), T(800)), (T(1e-10), T(1)))
+ elseif D == 3
+ limits!(ax, (T(8e-1), T(200)), (T(4e-5), T(1.5e0)))
+ end
+ path = "$output/prioranalysis"
+ ispath(path) || mkpath(path)
+ save("$path/spectra_$(D)D_dyadic_Re$(Int(Re)).pdf", fig)
+ fig
+end
+clean()
diff --git a/lib/PaperDC/src/PaperDC.jl b/lib/PaperDC/src/PaperDC.jl
new file mode 100644
index 000000000..b23d5b724
--- /dev/null
+++ b/lib/PaperDC/src/PaperDC.jl
@@ -0,0 +1,17 @@
+"""
+Utility functions for scripts.
+"""
+module PaperDC
+
+using IncompressibleNavierStokes
+using IncompressibleNavierStokes: momentum!, divergence!, project!, apply_bc_u!
+using Observables
+using LinearAlgebra
+
+include("observe.jl")
+include("rk.jl")
+
+export observe_u, observe_v
+export RKProject
+
+end # module PaperDC
diff --git a/lib/PaperDC/src/observe.jl b/lib/PaperDC/src/observe.jl
new file mode 100644
index 000000000..118037b69
--- /dev/null
+++ b/lib/PaperDC/src/observe.jl
@@ -0,0 +1,79 @@
+function observe_v(dnsobs, Φ, les, compression, psolver)
+ (; grid) = les
+ (; dimension, N, Iu, Ip) = grid
+ D = dimension()
+ Mα = N[1] - 2
+ v = zero.(Φ(dnsobs[].u, les, compression))
+ Pv = zero.(v)
+ p = zero(v[1])
+ div = zero(p)
+ ΦPF = zero.(v)
+ PFΦ = zero.(v)
+ c = zero.(v)
+ T = eltype(v[1])
+ results = (;
+ Φ,
+ Mα,
+ t = zeros(T, 0),
+ Dv = zeros(T, 0),
+ Pv = zeros(T, 0),
+ Pc = zeros(T, 0),
+ c = zeros(T, 0),
+ )
+ on(dnsobs) do (; u, PF, t)
+ push!(results.t, t)
+
+ Φ(v, u, les, compression)
+ apply_bc_u!(v, t, les)
+ Φ(ΦPF, PF, les, compression)
+ momentum!(PFΦ, v, nothing, t, les)
+ apply_bc_u!(PFΦ, t, les; dudt = true)
+ project!(PFΦ, les; psolver, div, p)
+ foreach(α -> c[α] .= ΦPF[α] .- PFΦ[α], 1:D)
+ apply_bc_u!(c, t, les)
+ divergence!(div, v, les)
+ norm_Du = norm(div[Ip])
+ norm_v = sqrt(sum(α -> sum(abs2, v[α][Iu[α]]), 1:D))
+ push!(results.Dv, norm_Du / norm_v)
+
+ copyto!.(Pv, v)
+ project!(Pv, les; psolver, div, p)
+ foreach(α -> Pv[α] .= Pv[α] .- v[α], 1:D)
+ norm_vmPv = sqrt(sum(α -> sum(abs2, Pv[α][Iu[α]]), 1:D))
+ push!(results.Pv, norm_vmPv / norm_v)
+
+ Pc = Pv
+ copyto!.(Pc, c)
+ project!(Pc, les; psolver, div, p)
+ foreach(α -> Pc[α] .= Pc[α] .- c[α], 1:D)
+ norm_cmPc = sqrt(sum(α -> sum(abs2, Pc[α][Iu[α]]), 1:D))
+ norm_c = sqrt(sum(α -> sum(abs2, c[α][Iu[α]]), 1:D))
+ push!(results.Pc, norm_cmPc / norm_c)
+
+ norm_ΦPF = sqrt(sum(α -> sum(abs2, ΦPF[α][Iu[α]]), 1:D))
+ push!(results.c, norm_c / norm_ΦPF)
+ end
+ results
+end
+
+observe_u(dns, psolver_dns, filters; nupdate = 1) =
+ processor() do state
+ PF = zero.(state[].u)
+ div = zero(state[].u[1])
+ p = zero(state[].u[1])
+ dnsobs = Observable((; state[].u, PF, state[].t))
+ results = [
+ observe_v(dnsobs, Φ, setup, compression, psolver) for
+ (; setup, Φ, compression, psolver) in filters
+ ]
+ on(state) do (; u, t, n)
+ n % nupdate == 0 || return
+ apply_bc_u!(u, t, dns)
+ momentum!(PF, u, nothing, t, dns)
+ apply_bc_u!(PF, t, dns; dudt = true)
+ project!(PF, dns; psolver = psolver_dns, div, p)
+ dnsobs[] = (; u, PF, t)
+ end
+ # state[] = state[] # Save initial conditions
+ results
+ end
diff --git a/lib/PaperDC/src/rk.jl b/lib/PaperDC/src/rk.jl
new file mode 100644
index 000000000..a2a944574
--- /dev/null
+++ b/lib/PaperDC/src/rk.jl
@@ -0,0 +1,149 @@
+"""
+ RKProject(rk, projectorder)
+
+Runge-Kutta method with different projection order.
+The Runge-Kutta method `rk` can be for example `RKMethods.RK44()`.
+
+- `projetorder = :first`: Project RHS before applying closure term.
+- `projetorder = :second`: Project RHS after applying closure term.
+- `projetorder = :last`: Project solution instead of RHS (same as `rk`).
+"""
+struct RKProject{T,R} <: IncompressibleNavierStokes.AbstractODEMethod{T}
+ rk::R
+ projectorder::Symbol
+ RKProject(rk, projectorder) = new{eltype(rk.A),typeof(rk)}(rk, projectorder)
+end
+
+ode_method_cache(method::RKProject, setup, u) = ode_method_cache(method.rk, setup, u)
+
+create_stepper(method::RKProject; setup, psolver, u, temp, t, n = 0) =
+ create_stepper(method.rk; setup, psolver, u, temp, t, n)
+
+function timestep!(method::RKProject, stepper, Δt; θ = nothing, cache)
+ (; setup, psolver, u, temp, t, n) = stepper
+ (; grid, closure_model) = setup
+ (; dimension, Iu) = grid
+ (; rk, projectorder) = method
+ (; A, b, c) = rk
+ (; u₀, ku, div, p) = cache
+ D = dimension()
+ nstage = length(b)
+ m = closure_model
+ projectorder ∈ (:first, :second, :last) || error("Unknown projectorder: $projectorder")
+
+ # Update current solution
+ t₀ = t
+ copyto!.(u₀, u)
+
+ for i = 1:nstage
+ # Compute force at current stage i
+ apply_bc_u!(u, t, setup)
+ momentum!(ku[i], u, temp, t, setup)
+
+ # Project F first
+ if projectorder == :first
+ apply_bc_u!(ku[i], t, setup; dudt = true)
+ project!(ku[i], setup; psolver, div, p)
+ end
+
+ # Add closure term
+ isnothing(m) || map((k, m) -> k .+= m, ku[i], m(u, θ))
+
+ # Project F second
+ if projectorder == :second
+ apply_bc_u!(ku[i], t, setup; dudt = true)
+ project!(ku[i], setup; psolver, div, p)
+ end
+
+ # Intermediate time step
+ t = t₀ + c[i] * Δt
+
+ # Apply stage forces
+ for α = 1:D
+ u[α] .= u₀[α]
+ for j = 1:i
+ @. u[α] += Δt * A[i, j] * ku[j][α]
+ # @. u[α][Iu[α]] += Δt * A[i, j] * ku[j][α][Iu[α]]
+ end
+ end
+
+ # Project stage u directly
+ # Make velocity divergence free at time t
+ if projectorder == :last
+ apply_bc_u!(u, t, setup)
+ project!(u, setup; psolver, div, p)
+ end
+ end
+
+ # This is redundant, but Neumann BC need to have _exact_ copies
+ # since we divide by an infinitely thin (eps(T)) volume width in the
+ # diffusion term
+ apply_bc_u!(u, t, setup)
+
+ create_stepper(method; setup, psolver, u, temp, t, n = n + 1)
+end
+
+function timestep(method::RKProject, stepper, Δt; θ = nothing)
+ (; setup, psolver, u, temp, t, n) = stepper
+ (; grid, closure_model) = setup
+ (; dimension) = grid
+ (; rk, projectorder) = method
+ (; A, b, c) = rk
+ D = dimension()
+ nstage = length(b)
+ m = closure_model
+ projectorder ∈ (:first, :second, :last) || error("Unknown projectorder: $projectorder")
+
+ # Update current solution (does not depend on previous step size)
+ t₀ = t
+ u₀ = u
+ ku = ()
+
+ for i = 1:nstage
+ # Compute force at current stage i
+ u = apply_bc_u(u, t, setup)
+ F = momentum(u, temp, t, setup)
+
+ # Project F first
+ if projectorder == :first
+ F = apply_bc_u(F, t, setup; dudt = true)
+ F = project(F, setup; psolver)
+ end
+
+ # Add closure term
+ isnothing(m) || (F = F .+ m(u, θ))
+
+ # Project F second
+ if projectorder == :second
+ F = apply_bc_u(F, t, setup; dudt = true)
+ F = project(F, setup; psolver)
+ end
+
+ # Store right-hand side of stage i
+ ku = (ku..., F)
+
+ # Intermediate time step
+ t = t₀ + c[i] * Δt
+
+ # Apply stage forces
+ u = u₀
+ for j = 1:i
+ u = @. u + Δt * A[i, j] * ku[j]
+ # u = tupleadd(u, @.(Δt * A[i, j] * ku[j]))
+ end
+
+ # Project stage u directly
+ # Make velocity divergence free at time t
+ if projectorder == :last
+ u = apply_bc_u(u, t, setup)
+ u = project(u, setup; psolver)
+ end
+ end
+
+ # This is redundant, but Neumann BC need to have _exact_ copies
+ # since we divide by an infinitely thin (eps(T)) volume width in the
+ # diffusion term
+ u = apply_bc_u(u, t, setup)
+
+ create_stepper(method; setup, psolver, u, temp, t, n = n + 1)
+end
diff --git a/lib/TensorClosure/Project.toml b/lib/TensorClosure/Project.toml
new file mode 100644
index 000000000..316021e0a
--- /dev/null
+++ b/lib/TensorClosure/Project.toml
@@ -0,0 +1,19 @@
+[deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+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"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
+LuxCUDA = "d0bbae9a-e099-4d5b-a835-1c6931763bda"
+NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+NeuralClosure = "099dac27-d7f2-4047-93d5-0baee36b9c25"
+Observables = "510215fc-4207-5dde-b226-833fc4488ee2"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
+Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
diff --git a/lib/TensorClosure/README.md b/lib/TensorClosure/README.md
new file mode 100644
index 000000000..1e4a0210b
--- /dev/null
+++ b/lib/TensorClosure/README.md
@@ -0,0 +1,27 @@
+# TensorClosure
+
+Tensor closure scripts.
+
+## Set up environment
+
+From this directory, run:
+
+```sh
+julia --project -e '
+using Pkg
+Pkg.develop([
+ PackageSpec(; path = "../.."),
+ PackageSpec(; path = "../NeuralClosure"),
+])
+Pkg.instantiate()
+'
+```
+
+or interactively from a Julia REPL:
+
+```julia-repl
+julia> ]
+(v1.10) pkg> activate .
+(TensorClosure) pkg> dev ../.. ../NeuralClosure
+(TensorClosure) pkg> instantiate
+```
diff --git a/lib/TensorClosure/main.jl b/lib/TensorClosure/main.jl
new file mode 100644
index 000000000..201fef4f6
--- /dev/null
+++ b/lib/TensorClosure/main.jl
@@ -0,0 +1,57 @@
+using CairoMakie
+using IncompressibleNavierStokes
+
+# For running on GPU
+using CUDA;
+CUDA.allowscalar(false);
+ArrayType = CuArray;
+T = Float32
+
+# Setup
+Re = T(10_000)
+n = 1024
+lims = T(0), T(1)
+x = LinRange(lims..., n + 1), LinRange(lims..., n + 1)
+setup = Setup(x...; Re, ArrayType);
+u₀ = random_field(setup, T(0));
+
+u = u₀
+
+B, V = IncompressibleNavierStokes.tensorbasis(u, setup)
+B |> length
+CUDA.@allowscalar B[2][5, 10]
+getindex.(B[2], 2)
+
+heatmap(Array(u[1]))
+u[1] |> Array |> heatmap
+getindex.(B[3], 1, 1) |> Array |> heatmap
+getindex.(B[3], 2, 1) |> Array |> heatmap
+getindex.(B[3], 2, 2) |> Array |> heatmap
+V[1] |> Array |> heatmap
+V[2] |> Array |> heatmap
+V[2] |> Array |> contourf
+
+# Solve unsteady problem
+state, outputs = solve_unsteady(
+ setup,
+ u₀,
+ (T(0), T(1));
+ Δt = T(1e-4),
+ processors = (
+ # rtp = realtimeplotter(; setup, nupdate = 1),
+ log = timelogger(; nupdate = 10),
+ ),
+);
+(; u) = state
+
+makeplot(u, fieldname, time) = save(
+ "output/fieldplots/$fieldname$time.png",
+ fieldplot((; u, t = T(0)); setup, fieldname, docolorbar = false, size = (500, 500)),
+)
+
+makeplot(u₀, :vorticity, 0)
+makeplot(u₀, :S2, 0)
+makeplot(u₀, :R2, 0)
+makeplot(u, :vorticity, 1)
+makeplot(u, :S2, 1)
+makeplot(u, :R2, 1)
diff --git a/scratch/Channel.jl b/scratch/Channel.jl
new file mode 100644
index 000000000..229377a41
--- /dev/null
+++ b/scratch/Channel.jl
@@ -0,0 +1,299 @@
+# Little LSP hack to get function signatures, go #src
+# to definition etc. #src
+if isdefined(@__MODULE__, :LanguageServer) #src
+ include("../src/IncompressibleNavierStokes.jl") #src
+ using .IncompressibleNavierStokes #src
+end #src
+
+using WGLMakie
+# using GLMakie
+# using CairoMakie
+using IncompressibleNavierStokes
+using FFTW
+using Adapt
+
+# Output directory
+output = "output/Channel"
+
+# Floating point type
+T = Float64
+
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+using CUDA;
+T = Float32;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+device = x -> adapt(CuArray{T}, x)
+
+set_theme!(; GLMakie = (; scalefactor = 1.5))
+
+# Reynolds number
+Re = T(4_000)
+
+# Boundary conditions
+boundary_conditions = (
+ ## x left, x right
+ (PeriodicBC(), PeriodicBC()),
+
+ ## y rear, y front
+ (DirichletBC(), DirichletBC()),
+)
+
+getdir(x = rand(T)) = T[cospi(2x), sinpi(2x)]
+createbodyforce((Ax, Ay), (a, b), r) = function bodyforce(dim, x, y, t)
+ A = Ax * (dim() == 1) + Ay * (dim() == 2)
+ R = (x - a)^2 / r^2 + (y - b)^2 / r^2
+ A * max(0, 1 - R)
+end
+createmanyforce(forces...) = function (dim, x, y, t)
+ out = zero(x)
+ for f in forces
+ out += f(dim, x, y, t)
+ end
+ out
+end
+
+bodyforce = createmanyforce(
+ createbodyforce(2 * getdir(+0.00), T[1.0, +0.0], T(0.6)),
+ createbodyforce(4 * getdir(-0.10), T[2.0, +0.5], T(0.2)),
+ createbodyforce(3 * getdir(+0.05), T[3.5, -0.3], T(0.3)),
+ createbodyforce(4 * getdir(-0.05), T[4.5, +0.1], T(0.2)),
+ createbodyforce(4 * getdir(-0.45), T[6.5, +0.5], T(0.2)),
+ createbodyforce(3 * getdir(+0.05), T[7.5, +0.0], T(0.3)),
+ createbodyforce(4 * getdir(-0.50), T[9.0, -0.5], T(0.2)),
+)
+
+bodyforce = createmanyforce(
+ createbodyforce(4 * getdir(+0.10), T[1.0, +0.4], T(0.3)),
+ createbodyforce(3 * getdir(-0.45), T[2.3, -0.1], T(0.2)),
+ createbodyforce(4 * getdir(+0.05), T[3.0, +0.2], T(0.2)),
+ createbodyforce(2 * getdir(-0.02), T[5.0, -0.1], T(0.5)),
+ createbodyforce(4 * getdir(-0.45), T[7.0, +0.4], T(0.2)),
+ createbodyforce(3 * getdir(+0.05), T[7.5, -0.5], T(0.3)),
+ createbodyforce(4 * getdir(+0.00), T[9.0, +0.0], T(0.2)),
+)
+
+# A 2D grid is a Cartesian product of two vectors. Here we refine the grid near
+# the walls.
+n = 64
+# n = 128
+x = LinRange(T(0), T(10), 8n + 1)
+y = LinRange(-T(1), T(1), 2n + 1)
+plotgrid(x, y)
+
+bodyforce.([IncompressibleNavierStokes.Dimension(1)], x, y', T(0)) |> heatmap
+bodyforce.([IncompressibleNavierStokes.Dimension(2)], x, y', T(0)) |> heatmap
+
+# Build setup and assemble operators
+setup = Setup(x, y; Re, bodyforce, boundary_conditions, ArrayType);
+
+psolver = DirectPressureSolver(setup);
+
+# Initial conditions
+u₀ = create_initial_conditions(
+ setup,
+ (dim, x, y) -> (dim() == 1) * 3 * (1 + y) * (1 - y);
+ # (dim, x, y) -> zero(x);
+ psolver,
+);
+u = u₀;
+
+# Solve unsteady problem
+state, outputs = solve_unsteady(
+ setup,
+ # u₀,
+ u,
+ (T(0), T(4.0));
+ Δt = T(0.005),
+ psolver,
+ processors = (
+ rtp = realtimeplotter(;
+ setup,
+ # type = contourf,
+ # plot = fieldplot,
+ # fieldname = :velocity,
+ # plot = energy_history_plot,
+ ## plot = energy_spectrum_plot,
+ nupdate = 1,
+ size = (1200, 600),
+ docolorbar = false,
+ ),
+ ## anim = animator(; setup, path = "$output/vorticity.mkv", nupdate = 20),
+ ## vtk = vtk_writer(; setup, nupdate = 10, dir = output, filename = "solution"),
+ ## field = fieldsaver(; setup, nupdate = 10),
+ log = timelogger(; nupdate = 1),
+ ),
+);
+(; u) = state;
+
+using CairoMakie
+GLMakie.activate!()
+CairoMakie.activate!()
+
+# Plot pressure
+fieldplot(
+ state;
+ setup,
+ psolver,
+ # type = contourf,
+ # fieldname = :pressure,
+ fieldname = :velocity,
+ # fieldname = :vorticity,
+ docolorbar = false,
+ # size = (1200, 350),
+ # size = (900, 250),
+ size = (650, 200),
+)
+
+name = "vorticity_10.pdf"
+name = "velocity_10.pdf"
+save(name, current_figure())
+run(`mv $name ../SupervisedClosure/figures/channel/`)
+
+(; Ip, Iu, xp) = setup.grid
+
+fig = with_theme(; palette = (; color = ["#3366cc", "#cc0000", "#669900", "#ffcc00"])) do
+ fig = Figure()
+ ax = Axis(fig[1, 1]; xlabel = "x")
+ lines!(ax, Array(xp[1][2:end-1]), u[2][Iu[2]][:, 32] |> Array)
+ fig
+end
+name = "midline_10.pdf"
+
+IncompressibleNavierStokes.kinetic_energy(u, setup)[Ip] |> Array |> heatmap
+sum(IncompressibleNavierStokes.kinetic_energy(u, setup)[Ip]; dims = 2) |> lines
+sum(IncompressibleNavierStokes.kinetic_energy(u, setup)[Ip]; dims = 1)[:] |> Array |> lines
+
+e = IncompressibleNavierStokes.kinetic_energy(u, setup)[Ip]
+e[:, 40]
+
+ehat = fft(e[:, n])[1:5n] |> x -> abs.(x) |> Array
+ehat2 = fft(e[:, 60])[1:5n] |> x -> abs.(x) |> Array
+lines(ehat; axis = (; xscale = log10, yscale = log10))
+lines!(ehat2)
+
+fig = Figure()
+Axis(fig[1, 1])
+u[2][Iu[2]][:, 32] |> Array |> lines!
+# u[2][Iu[2]][:, 40] |> Array |> lines!
+
+fig = Figure()
+Axis(fig[1, 1])
+u[1][Iu[1]][20, :] |> Array |> lines!
+u[1][Iu[1]][40, :] |> Array |> lines!
+u[1][Iu[1]][300, :] |> Array |> lines!
+
+# ## Post-process
+#
+# We may visualize or export the computed fields
+
+# Export to VTK
+save_vtk(setup, state.u, state.p, "$output/solution")
+
+Makie.available_gradients()
+
+#######################################################################
+
+bodyforce_train = createmanyforce(
+ createbodyforce(2 * getdir(+0.00), T[1.0, +0.0], T(0.6)),
+ createbodyforce(4 * getdir(-0.10), T[2.0, +0.5], T(0.2)),
+ createbodyforce(3 * getdir(+0.05), T[3.5, -0.3], T(0.3)),
+ createbodyforce(4 * getdir(-0.05), T[4.5, +0.1], T(0.2)),
+ createbodyforce(4 * getdir(-0.45), T[6.5, +0.5], T(0.2)),
+ createbodyforce(3 * getdir(+0.05), T[7.5, +0.0], T(0.3)),
+ createbodyforce(4 * getdir(-0.50), T[9.0, -0.5], T(0.2)),
+)
+
+bodyforce_test = createmanyforce(
+ createbodyforce(4 * getdir(+0.10), T[1.0, +0.4], T(0.3)),
+ createbodyforce(3 * getdir(-0.45), T[2.3, -0.1], T(0.2)),
+ createbodyforce(4 * getdir(+0.05), T[3.0, +0.2], T(0.2)),
+ createbodyforce(2 * getdir(-0.02), T[5.0, -0.1], T(0.5)),
+ createbodyforce(4 * getdir(-0.45), T[7.0, +0.4], T(0.2)),
+ createbodyforce(3 * getdir(+0.05), T[7.5, -0.5], T(0.3)),
+ createbodyforce(4 * getdir(+0.00), T[9.0, +0.0], T(0.2)),
+)
+
+# Parameters
+# nles = 50
+# nles = 64
+# nles = 128
+# ndns = 200
+nles = 32
+ndns = 128
+params = (;
+ D = 2,
+ Re = T(4_000),
+ lims = ((T(0), T(10)), (T(-1), T(1))),
+ nles = [(8nles, 2nles)],
+ ndns = (8ndns, 2ndns),
+ # tburn = T(0.1),
+ tsim = T(10),
+ Δt = T(1e-3),
+ savefreq = 5,
+ ArrayType,
+ boundary_conditions = ((PeriodicBC(), PeriodicBC()), (DirichletBC(), DirichletBC())),
+ PSolver = DirectPressureSolver,
+ icfunc = (setup, psolver) -> create_initial_conditions(
+ setup,
+ (dim, x, y) -> (dim() == 1) * 3 * (1 + y) * (1 - y);
+ # (dim, x, y) -> zero(x);
+ psolver,
+ ),
+)
+
+# Create LES data from DNS
+data_train = [create_les_data(T; params..., bodyforce = bodyforce_train) for _ = 1:1];
+data_test = [create_les_data(T; params..., bodyforce = bodyforce_test) for _ = 1:1];
+
+data_train[1].u[1][1][1]
+
+# Build LES setup and assemble operators
+x = ntuple(α -> LinRange(params.lims[α]..., params.nles[1][α] + 1), params.D)
+setup = Setup(x...; params.boundary_conditions, params.Re, ArrayType);
+
+# Uniform periodic grid
+psolver = params.PSolver(setup);
+
+# Inspect data
+field = data_train[1].u[1];
+
+u = device(field[1])
+# u = device(field[201])
+# u = device(field[1201])
+# u = device(field[2001])
+o = Observable((; u, t = nothing))
+fieldplot(
+ o;
+ setup,
+ # fieldname = :velocity,
+ # fieldname = 2,
+)
+
+# energy_spectrum_plot(o; setup)
+for i = 1:length(field)
+ o[] = (; o[]..., u = device(field[i]))
+ sleep(0.001)
+end
+
+io_train = create_io_arrays(data_train, [setup]);
+io_test = create_io_arrays(data_test, [setup]);
+
+closure, θ₀ = cnn(;
+ setup,
+ radii = [2, 2, 2, 2],
+ channels = [5, 5, 5, params.D],
+ activations = [leakyrelu, leakyrelu, leakyrelu, identity],
+ use_bias = [true, true, true, false],
+ rng,
+);
+closure.chain
+
+sample = io_train[1].u[:, :, :, 1:5]
+closure(sample, θ₀) |> size
diff --git a/scratch/energy.jl b/scratch/energy.jl
new file mode 100644
index 000000000..e37139b89
--- /dev/null
+++ b/scratch/energy.jl
@@ -0,0 +1,203 @@
+# Little LSP hack to get function signatures, go #src
+# to definition etc. #src
+if isdefined(@__MODULE__, :LanguageServer) #src
+ include("../src/IncompressibleNavierStokes.jl") #src
+ using .IncompressibleNavierStokes #src
+end #src
+
+# using GLMakie
+using CairoMakie
+using IncompressibleNavierStokes
+using IncompressibleNavierStokes: apply_bc_u!, total_kinetic_energy, diffusion!
+
+# Output directory
+output = "output/energy"
+mkdir(output)
+
+# Array type
+ArrayType = Array
+# using CUDA; ArrayType = CuArray;
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+clean() = GC.gc()
+
+using CUDA;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+clean() = (GC.gc(); CUDA.reclaim())
+
+set_theme!()
+set_theme!(; GLMakie = (; scalefactor = 1.5))
+
+# 2D
+T = Float64;
+Re = T(10_000)
+# ndns = 4096
+ndns = 1024
+# nles = 128
+# ndns = 256
+# nles = 128
+nles = 64
+compression = ndns ÷ nles
+D = 2
+kp = 20
+lims = T(0), T(1);
+
+dns, les = map((ndns, nles)) do n
+ setup = Setup(ntuple(α -> LinRange(lims..., n + 1), D)...; Re, ArrayType)
+ psolver = SpectralPressureSolver(setup)
+ (; setup, psolver)
+end;
+
+# Create random initial conditions
+u₀ = random_field(dns.setup, T(0); kp, dns.psolver);
+clean()
+
+state = (; u = u₀, t = T(0));
+
+observe_u(dns, les, compression; Δt, nupdate = 1) =
+ processor() do state
+ u = state[].u
+ v = zero.(FaceAverage()(u, les.setup, compression))
+ # u = copy.(u₀)
+ diffu = zero.(u)
+ diffv = zero.(v)
+ results = (;
+ t = zeros(T, 0),
+ Ku = zeros(T, 0),
+ Kv = zeros(T, 0),
+ Kuref = zeros(T, 0),
+ Kvref = zeros(T, 0),
+ )
+ on(state) do (; u, t, n)
+ n % nupdate == 0 || return
+ apply_bc_u!(u, t, dns.setup)
+ FaceAverage()(v, u, les.setup, compression)
+ apply_bc_u!(v, t, les.setup)
+ Ku = total_kinetic_energy(u, dns.setup)
+ Kv = total_kinetic_energy(v, les.setup)
+ push!(results.t, t)
+ push!(results.Ku, Ku)
+ push!(results.Kv, Kv)
+ if n == 0
+ push!(results.Kuref, Ku)
+ push!(results.Kvref, Kv)
+ else
+ fill!.(diffu, 0)
+ fill!.(diffv, 0)
+ diffusion!(diffu, u, dns.setup)
+ diffusion!(diffv, v, les.setup)
+ apply_bc_u!(diffu, t, dns.setup)
+ apply_bc_u!(diffv, t, les.setup)
+ for α = 1:D
+ diffu[α] .*= u[α]
+ diffv[α] .*= v[α]
+ end
+ sum(sum.(diffu)) / ndns^D
+ push!(
+ results.Kuref,
+ results.Kuref[end] + nupdate * Δt * sum(sum.(diffu)) / ndns^D,
+ )
+ push!(
+ results.Kvref,
+ results.Kvref[end] + nupdate * Δt * sum(sum.(diffv)) / nles^D,
+ )
+ end
+ end
+ state[] = state[] # Save initial conditions
+ results
+ end
+
+# Δt = 5e-5
+Δt = 1e-4
+
+# Solve unsteady problem
+@time state, outputs = solve_unsteady(
+ dns.setup,
+ u₀,
+ # state.u,
+ (T(0), T(5e-1));
+ Δt,
+ docopy = true,
+ dns.psolver,
+ processors = (
+ rtp = realtimeplotter(; dns.setup, displayupdates = true, nupdate = 50),
+ obs = observe_u(dns, les, compression; Δt, nupdate = 1),
+ log = timelogger(; nupdate = 1),
+ ),
+);
+clean()
+
+outputs.obs.t
+outputs.obs.Ku
+outputs.obs.Kuref
+outputs.obs.Kv
+outputs.obs.Kvref
+
+using CairoMakie
+
+fig = with_theme() do
+ fig = Figure()
+ ax = Axis(fig[1, 1]; xlabel = "t", title = "Kinetic energy")
+ lines!(
+ ax,
+ outputs.obs.t,
+ outputs.obs.Ku;
+ color = Cycled(1),
+ linestyle = :solid,
+ label = "DNS",
+ )
+ lines!(
+ ax,
+ outputs.obs.t,
+ outputs.obs.Kuref;
+ color = Cycled(1),
+ linestyle = :dash,
+ label = "DNS (reference)",
+ )
+ lines!(
+ ax,
+ outputs.obs.t,
+ outputs.obs.Kv;
+ color = Cycled(2),
+ linestyle = :solid,
+ label = "Filtered DNS",
+ )
+ lines!(
+ ax,
+ outputs.obs.t,
+ outputs.obs.Kvref;
+ color = Cycled(2),
+ linestyle = :dash,
+ label = "Filtered DNS (reference)",
+ )
+ axislegend(ax)
+ fig
+end
+
+save("$output/energy.pdf", fig)
+
+fig = with_theme() do
+ fig = Figure()
+ ax = Axis(fig[1, 1])
+ lines!(
+ ax,
+ outputs.obs.t,
+ outputs.obs.Kv ./ outputs.obs.Ku;
+ color = Cycled(2),
+ linestyle = :solid,
+ label = "Filtered DNS",
+ )
+ lines!(
+ ax,
+ outputs.obs.t,
+ outputs.obs.Kvref ./ outputs.obs.Ku;
+ color = Cycled(2),
+ linestyle = :dash,
+ label = "Filtered DNS (theoretical)",
+ )
+ # axislegend(ax)
+ fig
+end
diff --git a/scratch/filter_plot.jl b/scratch/filter_plot.jl
new file mode 100644
index 000000000..d380fccda
--- /dev/null
+++ b/scratch/filter_plot.jl
@@ -0,0 +1,151 @@
+# Little LSP hack to get function signatures, go #src
+# to definition etc. #src
+if isdefined(@__MODULE__, :LanguageServer) #src
+ include("../src/IncompressibleNavierStokes.jl") #src
+ using .IncompressibleNavierStokes #src
+end #src
+
+using GLMakie
+using IncompressibleNavierStokes
+
+# Floating point precision
+T = Float64
+
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+using CUDA;
+T = Float32;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+
+# Viscosity model
+Re = T(10_000)
+
+# A 2D grid is a Cartesian product of two vectors
+n = 1024
+lims = T(0), T(1)
+x = LinRange(lims..., n + 1), LinRange(lims..., n + 1)
+
+# Build setup and assemble operators
+setup = Setup(x...; Re, ArrayType);
+
+# Since the grid is uniform and identical for x and y, we may use a specialized
+# spectral pressure solver
+pressure_solver = SpectralPressureSolver(setup);
+
+u₀ = random_field(setup, T(0); pressure_solver);
+u = u₀
+
+function filter_plot(state, setup, les, comp; resolution = (1200, 600))
+ (; boundary_conditions, grid) = setup
+ (; dimension, xlims, x, xp, Ip) = grid
+ D = dimension()
+ xf = Array.(getindex.(setup.grid.xp, Ip.indices))
+ xfbar = Array.(getindex.(les.grid.xp, les.grid.Ip.indices))
+ (; u, t) = state[]
+ ω = IncompressibleNavierStokes.vorticity(u, setup)
+ ωp = IncompressibleNavierStokes.interpolate_ω_p(ω, setup)
+ _f = Array(ωp)[Ip]
+ f = @lift begin
+ sleep(0.001)
+ (; u, t) = $state
+ IncompressibleNavierStokes.apply_bc_u!(u, t, setup)
+ IncompressibleNavierStokes.vorticity!(ω, u, setup)
+ IncompressibleNavierStokes.interpolate_ω_p!(ωp, ω, setup)
+ copyto!(_f, view(ωp, Ip))
+ end
+ ubar = zero.(IncompressibleNavierStokes.face_average(u, les, comp))
+ ωbar = IncompressibleNavierStokes.vorticity(ubar, les)
+ ωpbar = IncompressibleNavierStokes.interpolate_ω_p(ωbar, les)
+ _g = Array(ωpbar)[les.grid.Ip]
+ g = @lift begin
+ (; u, t) = $state
+ IncompressibleNavierStokes.face_average!(ubar, u, les, comp)
+ IncompressibleNavierStokes.apply_bc_u!(ubar, t, les)
+ IncompressibleNavierStokes.vorticity!(ωbar, ubar, les)
+ IncompressibleNavierStokes.interpolate_ω_p!(ωpbar, ωbar, les)
+ copyto!(_g, view(ωpbar, les.grid.Ip))
+ end
+ lims = @lift IncompressibleNavierStokes.get_lims($f)
+ fig = Figure(; resolution)
+ ax, hm = heatmap(
+ fig[1, 1],
+ xf...,
+ f;
+ colorrange = lims,
+ axis = (;
+ title = "$(setup.grid.N .- 2)",
+ aspect = DataAspect(),
+ xlabel = "x",
+ ylabel = "y",
+ limits = (xlims[1]..., xlims[2]...),
+ ),
+ )
+ ax, hm = heatmap(
+ fig[1, 2],
+ xfbar...,
+ g;
+ colorrange = lims,
+ axis = (;
+ title = "$(les.grid.N .- 2)",
+ aspect = DataAspect(),
+ xlabel = "x",
+ # ylabel = "y",
+ # yticksvisible = false,
+ # yticklabelsvisible = false,
+ limits = (xlims[1]..., xlims[2]...),
+ ),
+ )
+ # Colorbar(fig[1, 3], hm)
+ display(fig)
+ fig
+end
+
+set_theme!(
+ theme_black();
+ # colormap = :plasma,
+ # colormap = :lajolla,
+ # colormap = :inferno,
+ # colormap = :hot,
+ # colormap = :magma,
+ # colormap = :Oranges,
+ # colormap = :Reds,
+ # colormap = :YlOrRd_9,
+ colormap = :seaborn_icefire_gradient,
+)
+
+comp = 8
+les = Setup(x[1][1:comp:end], x[2][1:comp:end]; Re, ArrayType)
+
+# Solve unsteady problem
+(; u)outputs = solve_unsteady(
+ setup,
+ u₀,
+ (T(0), T(2e0));
+ Δt = T(1e-4),
+ pressure_solver,
+ inplace = true,
+ processors = (
+ # realtimeplotter(; setup, nupdate = 10, docolorbar = false, displayupdates = false),
+ # animator(
+ # setup,
+ # "filtered.mp4";
+ # plotter = processor(state -> filter_plot(state, setup, les, comp; resolution = (1200, 600))),
+ # nupdate = 100,
+ # ),
+ processor(
+ state -> filter_plot(state, setup, les, comp; resolution = (1200, 600));
+ nupdate = 10,
+ ),
+ # energy_history_plotter(setup; nupdate = 20, displayfig = false),
+ # energy_spectrum_plotter(setup; nupdate = 10, displayfig = false),
+ ## vtk_writer(setup; nupdate = 10, dir = "output/$name", filename = "solution"),
+ ## fieldsaver(setup; nupdate = 10),
+ timelogger(; nupdate = 10),
+ ),
+);
diff --git a/scratch/multigrid.jl b/scratch/multigrid.jl
new file mode 100644
index 000000000..b3d33e3f9
--- /dev/null
+++ b/scratch/multigrid.jl
@@ -0,0 +1,661 @@
+# Little LSP hack to get function signatures, go #src
+# to definition etc. #src
+if isdefined(@__MODULE__, :LanguageServer) #src
+ include("../src/IncompressibleNavierStokes.jl") #src
+ using .IncompressibleNavierStokes #src
+end #src
+
+# # Train closure model
+#
+# Here, we consider a periodic box ``[0, 1]^2``. It is discretized with a
+# uniform Cartesian grid with square cells.
+
+using CairoMakie
+using GLMakie
+using IncompressibleNavierStokes
+using JLD2
+using LaTeXStrings
+using LinearAlgebra
+using Lux
+using NNlib
+using Optimisers
+using Random
+using Zygote
+using SparseArrays
+using KernelAbstractions
+using FFTW
+
+GLMakie.activate!()
+
+set_theme!(; GLMakie = (; scalefactor = 1.5))
+
+# Random number generator
+rng = Random.default_rng()
+Random.seed!(rng, 123)
+
+# Floating point precision
+T = Float64
+
+# Array type
+ArrayType = Array
+device = identity
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+using LuxCUDA
+using CUDA;
+T = Float32;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+device = cu
+
+# Parameters
+get_params(nles) = (;
+ D = 2,
+ Re = T(6_000),
+ lims = (T(0), T(1)),
+ tburn = T(0.05),
+ tsim = T(0.5),
+ Δt = T(1e-4),
+ nles = map(n -> (n, n), nles),
+ ndns = (2048, 2048),
+ ArrayType,
+ PSolver = SpectralPressureSolver,
+)
+
+params_train = (; get_params([64, 128, 256])..., savefreq = 5);
+params_valid = (; get_params([128])..., savefreq = 10);
+params_test = (; get_params([32, 64, 128, 256, 512, 1024])..., tsim = T(0.1), savefreq = 5);
+
+# Create LES data from DNS
+data_train = [create_les_data(T; params_train...) for _ = 1:5];
+data_valid = [create_les_data(T; params_valid...) for _ = 1:1];
+data_test = create_les_data(T; params_test...);
+
+# # Save filtered DNS data
+# jldsave("output/forced/data.jld2"; data_train, data_valid, data_test)
+
+# # Load previous LES data
+# data_train, data_valid, data_test = load("output/forced/data.jld2", "data_train", "data_valid", "data_test")
+
+# Build LES setup and assemble operators
+getsetups(params) = [
+ Setup(
+ ntuple(α -> LinRange(params.lims..., nles + 1), params.D)...;
+ params.Re,
+ params.ArrayType,
+ ) for nles in params.nles
+]
+setups_train = getsetups(params_train);
+setups_valid = getsetups(params_valid);
+setups_test = getsetups(params_test);
+
+# Create input/output arrays
+io_train = create_io_arrays(data_train, setups_train);
+io_valid = create_io_arrays(data_valid, setups_valid);
+
+# Inspect data
+ig = 3
+field, setup = data_train[1].u[ig], setups_train[ig];
+# field, setup = data_valid[1].u[ig], setups_valid[ig];
+# field, setup = data_test.u[ig], setups_test[ig];
+u = device(field[1]);
+o = Observable((; u, t = nothing));
+energy_spectrum_plot(o; setup)
+# fieldplot(
+# o;
+# setup,
+# # fieldname = :velocity,
+# # fieldname = 2,
+# )
+# energy_spectrum_plot( o; setup,)
+for i = 1:length(field)
+ o[] = (; o[]..., u = device(field[i]))
+ sleep(0.001)
+end
+
+GLMakie.activate!()
+CairoMakie.activate!()
+
+# Training data plot
+fig = with_theme() do
+ sample = data_train[1]
+ fig = Figure()
+ for (i, it) in enumerate((1, length(sample.t)))
+ for (j, ig) in enumerate((1, 2, 3))
+ setup = setups_train[ig]
+ xf = Array.(getindex.(setup.grid.xp, setup.grid.Ip.indices))
+ u = sample.u[ig][it] |> device
+ ωp =
+ IncompressibleNavierStokes.interpolate_ω_p(
+ IncompressibleNavierStokes.vorticity(u, setup),
+ setup,
+ )[setup.grid.Ip] |> Array
+ colorrange = IncompressibleNavierStokes.get_lims(ωp)
+ opts = (;
+ xticksvisible = false,
+ xticklabelsvisible = false,
+ yticklabelsvisible = false,
+ yticksvisible = false,
+ )
+ i == 2 && (
+ opts = (;
+ opts...,
+ xlabel = "x",
+ xticksvisible = true,
+ xticklabelsvisible = true,
+ )
+ )
+ j == 1 && (
+ opts = (;
+ opts...,
+ ylabel = "y",
+ yticklabelsvisible = true,
+ yticksvisible = true,
+ )
+ )
+ ax = Axis(
+ fig[i, j];
+ opts...,
+ title = "n = $(params_train.nles[ig]), t = $(round(sample.t[it]; digits = 1))",
+ aspect = DataAspect(),
+ limits = (params_test.lims..., params_test.lims...),
+ )
+ heatmap!(ax, xf..., ωp; colorrange)
+ end
+ end
+ fig
+end
+
+save("training_data.pdf", fig)
+
+# Plot training spectra
+fig = with_theme(; palette = (; color = ["#3366cc", "#cc0000", "#669900", "#ffcc00"])) do
+ ig = 3
+ setup = setups_train[ig]
+ (; xp, Ip) = setup.grid
+ D = params_train.D
+ sample = data_train[1]
+ K = size(Ip) .÷ 2
+ kx = ntuple(α -> 0:K[α]-1, D)
+ k = fill!(similar(xp[1], length.(kx)), 0)
+ for α = 1:D
+ kα = reshape(kx[α], ntuple(Returns(1), α - 1)..., :, ntuple(Returns(1), D - α)...)
+ k .+= kα .^ 2
+ end
+ k .= sqrt.(k)
+ k = reshape(k, :)
+ # Sum or average wavenumbers between k and k+1
+ nk = ceil(Int, maximum(k))
+ kmax = minimum(K) - 1
+ # kmax = nk - 1
+ kint = 1:kmax
+ ia = similar(xp[1], Int, 0)
+ ib = sortperm(k)
+ vals = similar(xp[1], 0)
+ ksort = k[ib]
+ jprev = 2 # Do not include constant mode
+ for ki = 1:kmax
+ j = findfirst(>(ki + 1), ksort)
+ isnothing(j) && (j = length(k) + 1)
+ ia = [ia; fill!(similar(ia, j - jprev), ki)]
+ # val = doaverage ? T(1) / (j - jprev) : T(1)
+ val = T(π) * ((ki + 1)^2 - ki^2) / (j - jprev)
+ vals = [vals; fill!(similar(vals, j - jprev), val)]
+ jprev = j
+ end
+ ib = ib[2:jprev-1]
+ A = sparse(ia, ib, vals, kmax, length(k))
+ # Build inertial slope above energy
+ krange = [T(kmax)^(T(2) / 3), T(kmax)]
+ # slope, slopelabel = D == 2 ? (-T(3), L"k^{-3}") : (-T(5 / 3), L"k^{-5/3}")
+ # slope, slopelabel = D == 2 ? (-T(3), "||k||₂⁻³") : (-T(5 / 3), "k⁻⁵³")
+ slope, slopelabel = D == 2 ? (-T(3), "|k|⁻³") : (-T(5 / 3), "|k|⁻⁵³")
+ slopeconst = T(0)
+ # Nice ticks
+ logmax = round(Int, log2(kmax + 1))
+ xticks = T(2) .^ (0:logmax)
+ # Make plot
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xticks,
+ # xlabel = "||k||₂",
+ xlabel = "|k|",
+ # xlabel = "k",
+ # xlabel = L"\| k \|_2",
+ # xlabel = L"|k|_2",
+ # ylabel = "e(k)",
+ ylabel = "e(|k|)",
+ # ylabel = L"e(\| k \|_2)",
+ # ylabel = L"e(|k|_2)",
+ title = "Kinetic energy (n = $(params_train.nles[ig]))",
+ xscale = log10,
+ yscale = log10,
+ limits = (extrema(kint)..., T(1e-8), T(1)),
+ )
+ for (i, it) in enumerate((1, length(sample.t)))
+ u = device.(sample.u[ig][it])
+ ke = IncompressibleNavierStokes.kinetic_energy(u, setup)
+ e = ke[Ip]
+ e = fft(e)[ntuple(α -> kx[α] .+ 1, D)...]
+ e = abs.(e) ./ size(e, 1)
+ e = A * reshape(e, :)
+ ehat = max.(e, eps(T)) # Avoid log(0)
+ slopeconst = max(slopeconst, maximum(ehat ./ kint .^ slope))
+ lines!(ax, kint, Array(ehat); label = "t = $(round(sample.t[it]; digits = 1))")
+ end
+ inertia = 2 .* slopeconst .* krange .^ slope
+ lines!(ax, krange, inertia; linestyle = :dash, label = slopelabel)
+ axislegend(ax; position = :lb)
+ autolimits!(ax)
+ fig
+end
+
+save("training_spectra.pdf", fig)
+
+closure, θ₀ = cnn(;
+ setup = setups_train[1],
+ radii = [2, 2, 2, 2],
+ channels = [20, 20, 20, params_train.D],
+ activations = [leakyrelu, leakyrelu, leakyrelu, identity],
+ use_bias = [true, true, true, false],
+ rng,
+);
+closure.chain
+
+# Prepare training
+loss = createloss(mean_squared_error, closure);
+dataloaders = createdataloader.(io_train; batchsize = 50, device);
+# dataloaders[1]()
+loss(dataloaders[1](), device(θ₀))
+it = rand(1:size(io_valid[1].u, 4), 50);
+validset = map(v -> v[:, :, :, it], io_valid[1]);
+
+# Prepare training
+θ = T(1.0e-1) * device(θ₀);
+opt = Optimisers.setup(Adam(T(1.0e-3)), θ);
+callbackstate = Point2f[];
+
+# Training with multiple grids at the same time
+(; opt, θ, callbackstate) = train(
+ dataloaders,
+ loss,
+ opt,
+ θ;
+ niter = 5000,
+ ncallback = 20,
+ callbackstate,
+ callback = create_callback(closure, device(validset)...; state = callbackstate),
+);
+GC.gc()
+CUDA.reclaim()
+
+# Extract parameters
+θ_cnn_shared = θ;
+
+# # Save trained parameters
+# jldsave("output/multigrid/theta_cnn_shared.jld2"; theta = Array(θ_cnn_shared))
+
+# # Load trained parameters
+# θθ = load("output/multigrid/theta_cnn_shared.jld2")
+# copyto!.(θ_cnn_shared, θθ["theta"])
+
+# Train grid-specialized closure models
+θ_cnn = map(dataloaders) do d
+ # Prepare training
+ θ = T(1.0e-1) * device(θ₀)
+ opt = Optimisers.setup(Adam(T(1.0e-3)), θ)
+ callbackstate = Point2f[]
+
+ # Training with multiple grids at the same time
+ (; opt, θ, callbackstate) = train(
+ [d],
+ loss,
+ opt,
+ θ;
+ niter = 5000,
+ ncallback = 20,
+ callbackstate,
+ callback = create_callback(closure, device(validset)...; state = callbackstate),
+ )
+ θ
+end
+GC.gc()
+CUDA.reclaim()
+
+# # Save trained parameters
+# jldsave("output/multigrid/theta_cnn.jld2"; theta = Array.(θ_cnn))
+
+# # Load trained parameters
+# θθ = load("output/multigrid/theta_cnn.jld2")
+# copyto!.(θ_cnn, θθ["theta"])
+
+# Train Smagorinsky closure model
+ig = 2;
+setup = setups_train[ig];
+sample = data_train[1];
+m = smagorinsky_closure(setup);
+θ = T(0.05)
+e_smag = sum(2:length(sample.t)) do it
+ It = setup.grid.Ip
+ u = sample.u[ig][it] |> device
+ c = sample.c[ig][it] |> device
+ mu = m(u, θ)
+ e = zero(eltype(u[1]))
+ for α = 1:D
+ # e += sum(abs2, mu[α][Ip] .- c[α][Ip]) / sum(abs2, c[α][Ip])
+ e += norm(mu[α][Ip] .- c[α][Ip]) / norm(c[α][Ip])
+ end
+ e / D
+end / length(sample.t)
+# for θ in LinRange(T(0), T(1), 100)];
+
+# lines(LinRange(T(0), T(1), 100), e_smag)
+
+# Errors for grid-specialized closure models
+e_cnn = zeros(T, length(θ_cnn))
+i_traintest = 2:4
+offset_i = 1
+for (i, setup) in enumerate(setups_test[2:4])
+ ig = i + offset_i
+ params = params_test[ig]
+ pressure_solver = SpectralPressureSolver(setup)
+ u = device.(data_test.u[ig])
+ u₀ = device(data_test.u[ig][1])
+ Δt = params_test.Δt * params_test.savefreq
+ tlims = extrema(data_test.t)
+ nupdate = 4
+ Δt /= nupdate
+ processors = (; relerr = relerr_trajectory(u, setup; nupdate))
+ closedsetup = (; setup..., closure_model = wrappedclosure(closure, θ_cnn[i], setup))
+ _, outputs = solve_unsteady(closedsetup, u₀, tlims; Δt, pressure_solver, processors)
+ e_cnn[i] = outputs.relerr[]
+end
+
+# Errors for all test grids
+e_nm = zeros(T, length(setups_test))
+e_smag = zeros(T, length(setups_test))
+e_cnn_shared = zeros(T, length(setups_test))
+for (ig, setup) in enumerate(setups_test)
+ @show ig
+ params = params_test[ig]
+ pressure_solver = SpectralPressureSolver(setup)
+ u = device.(data_test.u[ig])
+ u₀ = device(data_test.u[ig][1])
+ Δt = params_test.Δt * params_test.savefreq
+ tlims = extrema(data_test.t)
+ nupdate = 4
+ Δt /= nupdate
+ processors = (; relerr = relerr_trajectory(u, setup; nupdate))
+ _, outputs = solve_unsteady(setup, u₀, tlims; Δt, pressure_solver, processors)
+ e_nm[ig] = outputs.relerr[]
+ m = smagorinsky_closure(setup)
+ closedsetup = (; setup..., closure_model = u -> m(u, T(0.1)))
+ _, outputs = solve_unsteady(closedsetup, u₀, tlims; Δt, pressure_solver, processors)
+ e_smag[ig] = outputs.relerr[]
+ closedsetup = (; setup..., closure_model = wrappedclosure(closure, θ_cnn_shared, setup))
+ _, outputs = solve_unsteady(closedsetup, u₀, tlims; Δt, pressure_solver, processors)
+ e_cnn_shared[ig] = outputs.relerr[]
+end
+
+GC.gc()
+CUDA.reclaim()
+
+e_nm
+e_smag
+e_cnn
+e_cnn_shared
+# e_cnn = ones(T, length(setups_test))
+# e_fno_shared = ones(T, length(setups_test))
+# e_fno_spec = ones(T, length(setups_test))
+
+CairoMakie.activate!()
+GLMakie.activate!()
+
+# Plot convergence
+with_theme(;
+ # linewidth = 5,
+ # markersize = 10,
+ # markersize = 20,
+ # fontsize = 20,
+ palette = (; color = ["#3366cc", "#cc0000", "#669900", "#ffcc00"]),
+) do
+ nles = params_test.nles
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xscale = log10,
+ yscale = log10,
+ xticks = nles,
+ xlabel = "n",
+ title = "Relative error (DNS: n = $(params_test.ndns))",
+ )
+ scatterlines!(nles, e_nm; marker = :circle, label = "No closure")
+ scatterlines!(nles, e_smag; marker = :utriangle, label = "Smagorinsky")
+ scatterlines!(nles, e_cnn_shared; marker = :diamond, label = "CNN (shared)")
+ scatterlines!(params_train.nles, e_cnn; marker = :rect, label = "CNN (specialized)")
+ # scatterlines!(nles, e_fno_spec; label = "FNO (retrained)")
+ # scatterlines!(nles, e_fno_share; label = "FNO (shared parameters)")
+ lines!(
+ collect(extrema(nles[3:end])),
+ n -> 2e4 * n^-2.0;
+ linestyle = :dash,
+ label = "n⁻²",
+ )
+ axislegend(; position = :lb)
+ fig
+end
+
+save("convergence.pdf", current_figure())
+
+markers_labels = [
+ (:circle, ":circle"),
+ (:rect, ":rect"),
+ (:diamond, ":diamond"),
+ (:hexagon, ":hexagon"),
+ (:cross, ":cross"),
+ (:xcross, ":xcross"),
+ (:utriangle, ":utriangle"),
+ (:dtriangle, ":dtriangle"),
+ (:ltriangle, ":ltriangle"),
+ (:rtriangle, ":rtriangle"),
+ (:pentagon, ":pentagon"),
+ (:star4, ":star4"),
+ (:star5, ":star5"),
+ (:star6, ":star6"),
+ (:star8, ":star8"),
+ (:vline, ":vline"),
+ (:hline, ":hline"),
+ ('a', "'a'"),
+ ('B', "'B'"),
+ ('↑', "'\\uparrow'"),
+ ('😄', "'\\:smile:'"),
+ ('✈', "'\\:airplane:'"),
+]
+
+# Final spectra
+ig = 4
+setup = setups_test[ig];
+params = params_test
+pressure_solver = SpectralPressureSolver(setup);
+uref = device(data_test.u[ig][end]);
+u₀ = device(data_test.u[ig][1]);
+Δt = params_test.Δt * params_test.savefreq;
+tlims = extrema(data_test.t);
+nupdate = 4;
+Δt /= nupdate;
+state_nm, outputs = solve_unsteady(setup, u₀, tlims; Δt, pressure_solver);
+m = smagorinsky_closure(setup);
+closedsetup = (; setup..., closure_model = u -> m(u, T(0.1)));
+state_smag, outputs = solve_unsteady(closedsetup, u₀, tlims; Δt, pressure_solver);
+# closedsetup = (; setup..., closure_model = wrappedclosure(closure, θ_cnn_shared, setup));
+closedsetup = (; setup..., closure_model = wrappedclosure(closure, θ_cnn[ig-1], setup));
+state_cnn, outputs = solve_unsteady(closedsetup, u₀, tlims; Δt, pressure_solver);
+
+# Plot predicted spectra
+fig = with_theme(; palette = (; color = ["#3366cc", "#cc0000", "#669900", "#ffcc00"])) do
+ (; xp, Ip) = setup.grid
+ D = params.D
+ K = size(Ip) .÷ 2
+ kx = ntuple(α -> 0:K[α]-1, D)
+ k = fill!(similar(xp[1], length.(kx)), 0)
+ for α = 1:D
+ kα = reshape(kx[α], ntuple(Returns(1), α - 1)..., :, ntuple(Returns(1), D - α)...)
+ k .+= kα .^ 2
+ end
+ k .= sqrt.(k)
+ k = reshape(k, :)
+ # Sum or average wavenumbers between k and k+1
+ nk = ceil(Int, maximum(k))
+ kmax = nk - 1
+ # kmax = minimum(K) - 1
+ kint = 1:kmax
+ ia = similar(xp[1], Int, 0)
+ ib = sortperm(k)
+ vals = similar(xp[1], 0)
+ ksort = k[ib]
+ jprev = 2 # Do not include constant mode
+ for ki = 1:kmax
+ j = findfirst(>(ki + 1), ksort)
+ isnothing(j) && (j = length(k) + 1)
+ ia = [ia; fill!(similar(ia, j - jprev), ki)]
+ # val = doaverage ? T(1) / (j - jprev) : T(1)
+ val = T(π) * ((ki + 1)^2 - ki^2) / (j - jprev)
+ # val = T(1) / (j - jprev)
+ vals = [vals; fill!(similar(vals, j - jprev), val)]
+ jprev = j
+ end
+ ib = ib[2:jprev-1]
+ A = sparse(ia, ib, vals, kmax, length(k))
+ # Build inertial slope above energy
+ # krange = [cbrt(T(kmax)), T(kmax)]
+ krange = [T(kmax)^(T(2) / 3), T(kmax)]
+ # slope, slopelabel = D == 2 ? (-T(3), L"k^{-3}") : (-T(5 / 3), L"k^{-5/3}")
+ slope, slopelabel = D == 2 ? (-T(3), "|k|⁻³") : (-T(5 / 3), "|k|⁻⁵³")
+ slopeconst = T(0)
+ # Nice ticks
+ logmax = round(Int, log2(kmax + 1))
+ xticks = T(2) .^ (0:logmax)
+ # Make plot
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xticks,
+ xlabel = "|k|",
+ ylabel = "e(|k|)",
+ # title = "Kinetic energy (n = $(params.nles[ig])) at time t = $(round(data_test.t[end]; digits = 1))",
+ title = "Kinetic energy (n = $(params.nles[ig]))",
+ xscale = log10,
+ yscale = log10,
+ limits = (extrema(kint)..., T(1e-8), T(1)),
+ )
+ for (u, label) in (
+ # (uref, "Reference"),
+ (state_nm.u, "No closure"),
+ (state_smag.u, "Smagorinsky"),
+ (state_cnn.u, "CNN (specialized)"),
+ (uref, "Reference"),
+ )
+ ke = IncompressibleNavierStokes.kinetic_energy(u, setup)
+ e = ke[Ip]
+ e = fft(e)[ntuple(α -> kx[α] .+ 1, D)...]
+ e = abs.(e) ./ size(e, 1)
+ e = A * reshape(e, :)
+ ehat = max.(e, eps(T)) # Avoid log(0)
+ slopeconst = max(slopeconst, maximum(ehat ./ kint .^ slope))
+ lines!(ax, kint, Array(ehat); label)
+ end
+ inertia = 2 .* slopeconst .* krange .^ slope
+ lines!(ax, krange, inertia; linestyle = :dash, label = slopelabel)
+ axislegend(ax; position = :lb)
+ autolimits!(ax)
+ fig
+end
+
+save("predicted_spectra.pdf", fig)
+
+# Plot spectrum errors
+fig = with_theme(; palette = (; color = ["#3366cc", "#cc0000", "#669900", "#ffcc00"])) do
+ (; xp, Ip) = setup.grid
+ D = params.D
+ K = size(Ip) .÷ 2
+ kx = ntuple(α -> 0:K[α]-1, D)
+ k = fill!(similar(xp[1], length.(kx)), 0)
+ for α = 1:D
+ kα = reshape(kx[α], ntuple(Returns(1), α - 1)..., :, ntuple(Returns(1), D - α)...)
+ k .+= kα .^ 2
+ end
+ k .= sqrt.(k)
+ k = reshape(k, :)
+ # Sum or average wavenumbers between k and k+1
+ nk = ceil(Int, maximum(k))
+ # kmax = nk - 1
+ kmax = minimum(K) - 1
+ kint = 1:kmax
+ ia = similar(xp[1], Int, 0)
+ ib = sortperm(k)
+ vals = similar(xp[1], 0)
+ ksort = k[ib]
+ jprev = 2 # Do not include constant mode
+ for ki = 1:kmax
+ j = findfirst(>(ki + 1), ksort)
+ isnothing(j) && (j = length(k) + 1)
+ ia = [ia; fill!(similar(ia, j - jprev), ki)]
+ # val = doaverage ? T(1) / (j - jprev) : T(1)
+ val = T(π) * ((ki + 1)^2 - ki^2) / (j - jprev)
+ vals = [vals; fill!(similar(vals, j - jprev), val)]
+ jprev = j
+ end
+ ib = ib[2:jprev-1]
+ A = sparse(ia, ib, vals, kmax, length(k))
+ # Build inertial slope above energy
+ # krange = [cbrt(T(kmax)), T(kmax)]
+ krange = [T(kmax)^(T(2) / 3), T(kmax)]
+ # slope, slopelabel = D == 2 ? (-T(3), L"k^{-3}") : (-T(5 / 3), L"k^{-5/3}")
+ slope, slopelabel = D == 2 ? (-T(3), "|k|⁻³") : (-T(5 / 3), "|k|⁻⁵³")
+ slopeconst = T(0)
+ # Nice ticks
+ logmax = round(Int, log2(kmax + 1))
+ xticks = T(2) .^ (0:logmax)
+ # Make plot
+ fig = Figure(; size = (500, 400))
+ ax = Axis(
+ fig[1, 1];
+ xticks,
+ xlabel = "|k|",
+ ylabel = "e(|k|)",
+ # title = "Kinetic energy (n = $(params.nles[ig])) at time t = $(round(data_test.t[end]; digits = 1))",
+ title = "Relative energy error (n = $(params.nles[ig]))",
+ xscale = log10,
+ yscale = log10,
+ limits = (extrema(kint)..., T(1e-8), T(1)),
+ )
+ ke = IncompressibleNavierStokes.kinetic_energy(uref, setup)
+ e = ke[Ip]
+ e = fft(e)[ntuple(α -> kx[α] .+ 1, D)...]
+ e = abs.(e) ./ size(e, 1)
+ e = A * reshape(e, :)
+ eref = max.(e, eps(T)) # Avoid log(0)
+ for (u, label) in (
+ (state_nm.u, "No closure"),
+ (state_smag.u, "Smagorinsky"),
+ (state_cnn.u, "CNN (specialized)"),
+ )
+ ke = IncompressibleNavierStokes.kinetic_energy(u, setup)
+ e = ke[Ip]
+ e = fft(e)[ntuple(α -> kx[α] .+ 1, D)...]
+ e = abs.(e) ./ size(e, 1)
+ e = A * reshape(e, :)
+ ehat = max.(e, eps(T)) # Avoid log(0)
+ ee = @. abs(ehat - eref) / abs(eref)
+ lines!(ax, kint, Array(ee); label)
+ end
+ axislegend(ax; position = :lt)
+ autolimits!(ax)
+ fig
+end
+
+save("spectrum_error.pdf", fig)
diff --git a/scratch/testgrad.jl b/scratch/testgrad.jl
new file mode 100644
index 000000000..bf31744d8
--- /dev/null
+++ b/scratch/testgrad.jl
@@ -0,0 +1,192 @@
+# Little LSP hack to get function signatures, go #src
+# to definition etc. #src
+if isdefined(@__MODULE__, :LanguageServer) #src
+ include("../src/IncompressibleNavierStokes.jl") #src
+ using .IncompressibleNavierStokes #src
+end #src
+
+#md using CairoMakie
+using GLMakie #!md
+using IncompressibleNavierStokes
+using KernelAbstractions
+using Zygote
+using LinearAlgebra
+using Random
+
+set_theme!(; GLMakie = (; scalefactor = 1.5))
+
+# Floating point precision
+T = Float64
+
+# Array type
+ArrayType = Array
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+using CUDA;
+T = Float32;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+
+# Viscosity model
+Re = T(100)
+
+# A 2D grid is a Cartesian product of two vectors
+# n = 8
+n = 16
+# n = 32
+# n = 128
+# n = 1024
+# n = 2056
+lims = T(0), T(1)
+x = LinRange(lims..., n + 1), LinRange(lims..., n + 1)
+# plotgrid(x...)
+
+# Build setup and assemble operators
+setup = Setup(x...; Re, ArrayType);
+
+# Since the grid is uniform and identical for x and y, we may use a specialized
+# spectral pressure solver
+psolver = SpectralPressureSolver(setup);
+
+u₀ = random_field(setup, T(0); psolver);
+u = u₀
+
+(; Iu, Ip) = setup.grid
+
+function finitediff(f, u::Tuple, α, I; h = sqrt(eps(eltype(u[1]))))
+ u1 = copy.(u)
+ CUDA.@allowscalar u1[α][I] -= h / 2
+ r1 = f(u1)
+ u2 = copy.(u)
+ CUDA.@allowscalar u2[α][I] += h / 2
+ r2 = f(u2)
+ (r2 - r1) / h
+end
+
+function finitediff(f, p, I; h = sqrt(eps(eltype(p))))
+ p1 = copy(p)
+ CUDA.@allowscalar p1[I] -= h / 2
+ r1 = f(p1)
+ p2 = copy(p)
+ CUDA.@allowscalar p2[I] += h / 2
+ r2 = f(p2)
+ (r2 - r1) / h
+end
+
+IncompressibleNavierStokes.diffusion(u, setup)[1]
+IncompressibleNavierStokes.diffusion_adjoint!(zero.(u), u, setup)[1]
+
+gradient(u -> sum(IncompressibleNavierStokes.diffusion(u, setup)[1]), u)[1][1]
+
+# Divergence
+ur = randn!.(similar.(u))
+φ = IncompressibleNavierStokes.divergence(ur, setup)
+φbar = randn!(similar(φ))
+dot(φbar, φ)
+dot(IncompressibleNavierStokes.divergence_adjoint!(zero.(ur), φbar, setup), ur)
+
+# Diffusion
+φ = IncompressibleNavierStokes.diffusion(u, setup)
+φbar = randn!.(similar.(φ))
+ubar = IncompressibleNavierStokes.diffusion_adjoint!(zero.(u), φbar, setup)
+dot(φbar, φ)
+dot(ubar, u)
+
+# Convection
+ur = randn!.(similar.(u))
+# ur = zero.(u)
+# for α = 1:length(ur)
+# ur[α][Iu[α]] .= randn.()
+# end
+# ur = u
+φ = IncompressibleNavierStokes.convection!(zero.(ur), ur, setup)
+φbar = randn!.(similar.(φ))
+ubar = IncompressibleNavierStokes.convection_adjoint!(zero.(ur), φbar, ur, setup)
+dot(φbar, φ)
+dot(ubar, ur) / 2
+
+function f(u, setup)
+ (; Iu) = setup.grid
+ u = IncompressibleNavierStokes.apply_bc_u(u, T(0), setup)
+ φ = IncompressibleNavierStokes.momentum(u, T(0), setup)
+ # φ = IncompressibleNavierStokes.diffusion(u, setup)
+ # φ = IncompressibleNavierStokes.convection(u, setup)
+ # φ = u
+ # dot(φ, φ)
+ dot(getindex.(φ, Iu), getindex.(φ, Iu))
+ # sum(abs2, getindex.(φ, Iu))
+ # u[1][1]
+end
+
+I = CartesianIndex(1, 2)
+α = 1
+CUDA.@allowscalar gradient(u -> f(u, setup), u)[1][α][I]
+finitediff(u -> f(u, setup), u, α, I)
+
+function f(u, setup)
+ (; Ω, Iu) = setup.grid
+ φ = u
+ φ = IncompressibleNavierStokes.apply_bc_u(u, T(0), setup)
+ φ = IncompressibleNavierStokes.momentum(φ, T(0), setup)
+ φ = IncompressibleNavierStokes.apply_bc_u(φ, T(0), setup)
+ φ = IncompressibleNavierStokes.project(φ, setup; psolver)
+ # dot(φ, φ)
+ dot(getindex.(φ, Iu), getindex.(φ, Iu))
+ # sum(abs2, getindex.(φ, Iu))
+end
+
+I = CartesianIndex(8, 2)
+α = 1
+CUDA.@allowscalar gradient(u -> f(u, setup), u)[1][α][I]
+finitediff(u -> f(u, setup), u, α, I)
+
+function f(u, setup)
+ (; Ω, Iu) = setup.grid
+ method = RK44(; T)
+ stepper = IncompressibleNavierStokes.create_stepper(method; setup, psolver, u, t = T(0))
+ stepper = IncompressibleNavierStokes.timestep(method, stepper, T(1e-4))
+ φ = stepper.u
+ # dot(φ, φ)
+ dot(getindex.(φ, Iu), getindex.(φ, Iu))
+ # sum(abs2, getindex.(φ, Iu))
+end
+
+I = CartesianIndex(5, 6)
+α = 1
+CUDA.@allowscalar gradient(u -> f(u, setup), u)[1][α][I]
+finitediff(u -> f(u, setup), u, α, I)
+
+function f(u, setup)
+ (; Iu) = setup.grid
+ φ = IncompressibleNavierStokes.tupleadd(u, u)
+ dot(getindex.(φ, Iu), getindex.(φ, Iu))
+end
+
+IncompressibleNavierStokes.tupleadd(u, u)
+
+f(u, setup)
+
+I = CartesianIndex(2, 2)
+CUDA.@allowscalar gradient(u -> f(u, setup), u)[1][1][I]
+finitediff(u -> f(u, setup), u, I)
+
+function fp(p, setup)
+ (; Ip) = setup.grid
+ p = IncompressibleNavierStokes.apply_bc_p(p, T(0), setup)
+ φ = IncompressibleNavierStokes.pressuregradient(p, setup)
+ dot(φ, φ)
+ # dot(getindex.(φ, Iu), getindex.(φ, Iu))
+ # sum(abs2, getindex.(φ, Iu))
+end
+
+I = CartesianIndex(2, 2)
+CUDA.@allowscalar gradient(p -> fp(p, setup), p)[1][I]
+finitediff(p -> fp(p, setup), p, I)
+
+φ[1]
+φbar[1]
+ur[1]
+ubar[1]
diff --git a/scratch/train_model.jl b/scratch/train_model.jl
new file mode 100644
index 000000000..15c862ab4
--- /dev/null
+++ b/scratch/train_model.jl
@@ -0,0 +1,365 @@
+# Little LSP hack to get function signatures, go #src
+# to definition etc. #src
+if isdefined(@__MODULE__, :LanguageServer) #src
+ include("../src/IncompressibleNavierStokes.jl") #src
+ using .IncompressibleNavierStokes #src
+end #src
+
+# # Train closure model
+#
+# Here, we consider a periodic box ``[0, 1]^2``. It is discretized with a
+# uniform Cartesian grid with square cells.
+
+using Adapt
+using GLMakie
+using IncompressibleNavierStokes
+using JLD2
+using LinearAlgebra
+using Lux
+using NNlib
+using Optimisers
+using Random
+using Zygote
+
+# Random number generator
+rng = Random.default_rng()
+Random.seed!(rng, 123)
+
+set_theme!(; GLMakie = (; scalefactor = 1.5))
+
+# Floating point precision
+T = Float64
+
+# Array type
+ArrayType = Array
+device = identity
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+using LuxCUDA
+using CUDA;
+T = Float32;
+# T = Float64;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+# device = cu
+device = x -> adapt(CuArray{T}, x)
+
+# Parameters
+# nles = 50
+# nles = 64
+# nles = 128
+# ndns = 200
+nles = 256
+params = (;
+ D = 2,
+ Re = T(6_000),
+ lims = (T(0), T(1)),
+ nles = [(nles, nles)],
+ # ndns = (512, 512),
+ ndns = (1024, 1024),
+ tburn = T(0.05),
+ tsim = T(0.5),
+ Δt = T(1e-4),
+ savefreq = 5,
+ ArrayType,
+ PSolver = SpectralPressureSolver,
+)
+
+# Create LES data from DNS
+data_train = [create_les_data(T; params...) for _ = 1:3];
+data_valid = [create_les_data(T; params...) for _ = 1:1];
+data_test = [create_les_data(T; params..., tsim = T(0.2)) for _ = 1:1];
+
+data_test[1].u[1][1][1]
+
+# # Save filtered DNS data
+# jldsave("output/forced/data.jld2"; data_train, data_valid, data_test)
+
+# # Load previous LES data
+# data_train, data_valid, data_test = load("output/forced/data.jld2", "data_train", "data_valid", "data_test")
+
+# Build LES setup and assemble operators
+x = ntuple(α -> LinRange(params.lims..., nles + 1), params.D)
+setup = Setup(x...; params.Re, ArrayType);
+
+# Uniform periodic grid
+psolver = SpectralPressureSolver(setup);
+
+# Inspect data
+(; Ip) = setup.grid;
+field = data_train[1].u[1];
+α = 2
+# j = 13
+o = Observable(field[1][α][Ip])
+# o = Observable(field[1][α][:, :, j])
+heatmap(o)
+for i = 1:length(field)
+ o[] = field[i][α][Ip]
+ # o[] = field[i][α][:, :, j]
+ sleep(0.001)
+end
+
+# Inspect data
+field = data_train[1].u[1];
+u = device(field[1])
+o = Observable((; u, t = nothing))
+fieldplot(
+ o;
+ setup,
+ # fieldname = :velocity,
+ # fieldname = 2,
+)
+# energy_spectrum_plot(o; setup)
+for i = 1:length(field)
+ o[] = (; o[]..., u = device(field[i]))
+ sleep(0.001)
+end
+
+# Create input/output arrays
+io_train = create_io_arrays(data_train, [setup]);
+io_valid = create_io_arrays(data_valid, [setup]);
+io_test = create_io_arrays(data_test, [setup]);
+
+size(io_train[1].u)
+size(io_valid[1].u)
+size(io_test[1].u)
+
+Base.summarysize(io_train) / 1e9
+
+closure, θ₀ = cnn(;
+ setup,
+ radii = [2, 2, 2, 2],
+ channels = [5, 5, 5, params.D],
+ activations = [leakyrelu, leakyrelu, leakyrelu, identity],
+ use_bias = [true, true, true, false],
+ rng,
+);
+closure.chain
+
+sample = io_train[1].u[:, :, :, 1:5]
+closure(sample, θ₀) |> size
+
+θ₀.layer_5
+θ₀.layer_6
+
+# closure, θ₀ = fno(;
+# setup,
+# kmax = [20, 20, 20, 20],
+# c = [24, 12, 8, 8],
+# # σ = [gelu, gelu, gelu, identity],
+# σ = [leakyrelu, leakyrelu, leakyrelu, identity],
+# # σ = [tanh, tanh, tanh, identity],
+# # ψ = gelu,
+# ψ = tanh,
+# rng,
+# );
+#
+# closure.chain
+#
+# θ₀.layer_4.spectral_weights |> size
+
+# θ₀
+# θ₀ = 2*θ₀
+# θ₀.layer_6 ./= 20
+
+# Prepare training
+θ = T(1.0e-1) * device(θ₀);
+# θ = device(θ₀);
+# θ = 2 * device(θ₀);
+opt = Optimisers.setup(Adam(T(1.0e-3)), θ);
+callbackstate = Point2f[];
+it = rand(1:size(io_valid[1].u, 4), 50);
+validset = map(v -> v[:, :, :, it], io_valid[1]);
+
+# A-priori loss
+loss = createloss(mean_squared_error, closure);
+dataloader = createdataloader(io_train[1]; batchsize = 50, device);
+dataloaders = [dataloader]
+dataloader()
+
+# A-posteriori loss
+loss = IncompressibleNavierStokes.create_trajectory_loss(; setup, psolver, closure);
+dataloaders = [
+ IncompressibleNavierStokes.createtrajectoryloader(data_train; device, nunroll = 20)
+ for _ = 1:4
+];
+loss(dataloaders[1](), device(θ₀))
+
+# Warm-up
+loss(dataloaders[1](), θ)
+@time loss(dataloaders[1](), θ);
+b = dataloaders[1]();
+first(gradient(θ -> loss(b, θ), θ));
+@time first(gradient(θ -> loss(b, θ), θ));
+GC.gc()
+CUDA.reclaim()
+
+map() do
+ i = 3
+ # h = 1000 * sqrt(eps(T))
+ h = cbrt(eps(T))
+ θ1 = copy(θ)
+ θ2 = copy(θ)
+ CUDA.@allowscalar θ1[i] -= h / 2
+ CUDA.@allowscalar θ2[i] += h / 2
+ b = dataloaders[1]()
+ @show loss(b, θ2) loss(b, θ1)
+ a = (loss(b, θ2) - loss(b, θ1)) / h
+ b = CUDA.@allowscalar first(gradient(θ -> loss(b, θ), θ))[i]
+ [a; b]
+end
+
+# Training
+# Note: The states `opt`, `θ`, and `callbackstate`
+# will not be overwritten until training is finished.
+# This allows for cancelling with "Control-C" should errors explode.
+(; opt, θ, callbackstate) = train(
+ dataloaders,
+ loss,
+ opt,
+ θ;
+ niter = 1000,
+ ncallback = 10,
+ callbackstate,
+ callback = create_callback(closure, device(validset)...; state = callbackstate),
+);
+GC.gc()
+CUDA.reclaim()
+
+Array(θ)
+
+# # Save trained parameters
+# jldsave("output/forced/theta_cnn.jld2"; theta = Array(θ))
+# jldsave("output/forced/theta_fno.jld2"; theta = Array(θ))
+
+# # Load trained parameters
+# θθ = load("output/theta_cnn.jld2")
+# θθ = load("output/theta_fno.jld2")
+# copyto!(θ, θθ["theta"])
+
+function relerr(u, uref, setup)
+ (; dimension, Ip) = setup.grid
+ D = dimension()
+ a, b = T(0), T(0)
+ for α = 1:D
+ a += sum(abs2, u[α][Ip] - uref[α][Ip])
+ b += sum(abs2, uref[α][Ip])
+ end
+ sqrt(a) / sqrt(b)
+end
+
+u, u₀, = nothing, nothing
+u = device.(data_test[1].u[1]);
+u₀ = device(data_test[1].u[1][1]);
+Δt = data_test[1].t[2] - data_test[1].t[1]
+tlims = extrema(data_test[1].t)
+length(u)
+length(data_test[1].t)
+
+state_nm, outputs = solve_unsteady(
+ setup,
+ u₀,
+ tlims;
+ Δt,
+ psolver,
+ processors = (;
+ relerr = relerr_trajectory(u, setup),
+ log = timelogger(; nupdate = 1000),
+ ),
+)
+relerr_nm = outputs.relerr[]
+
+state_cnn, outputs = solve_unsteady(
+ (; setup..., closure_model = wrappedclosure(closure, θ, setup)),
+ u₀,
+ tlims;
+ Δt,
+ psolver,
+ processors = (relerr = relerr_trajectory(u, setup), log = timelogger(; nupdate = 1)),
+)
+relerr_cnn = outputs.relerr[]
+
+relerr_nm
+relerr_cnn
+
+# dnm = relerr_nm
+# dcnn = relerr_cnn
+
+dnm
+dcnn
+
+function energy_history(setup, state)
+ (; Ωp) = setup.grid
+ points = Point2f[]
+ on(state) do (; V, t)
+ V = Array(V)
+ vels = get_velocity(setup, V, t)
+ vels = reshape.(vels, :)
+ E = sum(vel -> sum(@. Ωp * vel^2), vels)
+ push!(points, Point2f(t, E))
+ end
+ points
+end
+
+energy_history_writer(setup; nupdate = 1, kwargs...) =
+ processor(state -> energy_history(setup, state; kwargs...); nupdate)
+
+isample = 1
+forcedsetup = (; setup..., force = data_train.force[:, isample]);
+
+devsetup = device(forcedsetup);
+V_nm, outputs_nm = solve_unsteady(
+ forcedsetup,
+ data_test.V[:, 1, isample],
+ (T(0), tsim);
+ Δt = T(2e-4),
+ processors = (
+ field_plotter(devsetup; type = heatmap, nupdate = 1),
+ energy_history_writer(forcedsetup),
+ step_logger(; nupdate = 10),
+ ),
+ psolver,
+ inplace = false,
+ device,
+ devsetup,
+)
+ehist_nm = outputs_nm[2]
+
+box = [
+ Point2f(0.72, 0.42),
+ Point2f(0.81, 0.42),
+ Point2f(0.81, 0.51),
+ Point2f(0.72, 0.51),
+ Point2f(0.72, 0.42),
+]
+
+plot_vorticity(setup, V, tsim)
+lines!(box; color = :red)
+current_figure()
+
+save("output/train/vorticity.png", current_figure())
+
+plot_vorticity(setup, V_nm, tsim)
+lines!(box; color = :red)
+current_figure()
+
+save("output/train/vorticity_nm.png", current_figure())
+
+plot_vorticity(setup, V_fno, tsim)
+lines!(box; color = :red)
+current_figure()
+
+save("output/train/vorticity_fno.png", current_figure())
+
+heatmap(vcat(
+ selectdim(reshape(V_nm, n, n, 2), 3, 1),
+ # selectdim(reshape(V_fno, n, n, 2), 3, 1),
+ selectdim(reshape(V, n, n, 2), 3, 1),
+))
+
+CUDA.memory_status()
+GC.gc()
+CUDA.reclaim()
diff --git a/scratch/trajectoryfit.jl b/scratch/trajectoryfit.jl
new file mode 100644
index 000000000..0dd117376
--- /dev/null
+++ b/scratch/trajectoryfit.jl
@@ -0,0 +1,292 @@
+# Little LSP hack to get function signatures, go #src
+# to definition etc. #src
+if isdefined(@__MODULE__, :LanguageServer) #src
+ include("../src/IncompressibleNavierStokes.jl") #src
+ using .IncompressibleNavierStokes #src
+end #src
+
+# # Train closure model
+#
+# Here, we consider a periodic box ``[0, 1]^2``. It is discretized with a
+# uniform Cartesian grid with square cells.
+
+using Adapt
+using GLMakie
+using IncompressibleNavierStokes
+using JLD2
+using LinearAlgebra
+using Lux
+using NNlib
+using Optimisers
+using Random
+using Zygote
+
+# Random number generator
+rng = Random.default_rng()
+Random.seed!(rng, 123)
+
+set_theme!(; GLMakie = (; scalefactor = 1.5))
+
+# Floating point precision
+T = Float64
+
+# Array type
+ArrayType = Array
+device = identity
+## using CUDA; ArrayType = CuArray
+## using AMDGPU; ArrayType = ROCArray
+## using oneAPI; ArrayType = oneArray
+## using Metal; ArrayType = MtlArray
+
+using LuxCUDA
+using CUDA;
+T = Float32;
+# T = Float64;
+ArrayType = CuArray;
+CUDA.allowscalar(false);
+# device = cu
+device = x -> adapt(CuArray{T}, x)
+
+# Parameters
+# nles = 50
+# nles = 64
+# nles = 128
+# ndns = 200
+nles = 256
+params = (;
+ D = 2,
+ Re = T(6_000),
+ lims = (T(0), T(1)),
+ nles = [nles],
+ # ndns = 512,
+ # ndns = 1024,
+ ndns = 2048,
+ tburn = T(0.05),
+ tsim = T(0.5),
+ Δt = T(1e-4),
+ savefreq = 5,
+ ArrayType,
+)
+
+# Create LES data from DNS
+data_train = [create_les_data(T; params...) for _ = 1:3];
+data_valid = [create_les_data(T; params...) for _ = 1:1];
+data_test = [create_les_data(T; params..., tsim = T(0.2)) for _ = 1:1];
+
+data_test[1].u[1][1][1]
+
+# # Save filtered DNS data
+# jldsave("output/forced/data.jld2"; data_train, data_valid, data_test)
+
+# # Load previous LES data
+# data_train, data_valid, data_test = load("output/forced/data.jld2", "data_train", "data_valid", "data_test")
+
+# Build LES setup and assemble operators
+x = ntuple(α -> LinRange(params.lims..., nles + 1), params.D)
+setup = Setup(x...; params.Re, ArrayType);
+
+# Uniform periodic grid
+pressure_solver = SpectralPressureSolver(setup);
+
+# Inspect data
+(; Ip) = setup.grid;
+field = data_train[1].u[1];
+α = 2
+# j = 13
+o = Observable(field[1][α][Ip])
+# o = Observable(field[1][α][:, :, j])
+heatmap(o)
+for i = 1:length(field)
+ o[] = field[i][α][Ip]
+ # o[] = field[i][α][:, :, j]
+ sleep(0.001)
+end
+
+# Inspect data
+field = data_train[1].u[1];
+u = device(field[1])
+o = Observable((; u, t = nothing))
+fieldplot(
+ o;
+ setup,
+ # fieldname = :velocity,
+ # fieldname = 2,
+)
+# energy_spectrum_plot(o; setup)
+for i = 1:length(field)
+ o[] = (; o[]..., u = device(field[i]))
+ sleep(0.001)
+end
+
+# Create input/output arrays
+io_train = create_io_arrays(data_train, [setup]);
+io_valid = create_io_arrays(data_valid, [setup]);
+io_test = create_io_arrays(data_test, [setup]);
+
+size(io_train[1].u)
+size(io_valid[1].u)
+size(io_test[1].u)
+
+Base.summarysize(io_train) / 1e9
+
+closure, θ₀ = cnn(;
+ setup,
+ radii = [2, 2, 2, 2],
+ channels = [5, 5, 5, params.D],
+ activations = [leakyrelu, leakyrelu, leakyrelu, identity],
+ use_bias = [true, true, true, false],
+ rng,
+);
+closure.chain
+
+sample = io_train[1].u[:, :, :, 1:5]
+closure(sample, θ₀) |> size
+
+θ₀.layer_5
+θ₀.layer_6
+
+# closure, θ₀ = fno(;
+# setup,
+# kmax = [20, 20, 20, 20],
+# c = [24, 12, 8, 8],
+# # σ = [gelu, gelu, gelu, identity],
+# σ = [leakyrelu, leakyrelu, leakyrelu, identity],
+# # σ = [tanh, tanh, tanh, identity],
+# # ψ = gelu,
+# ψ = tanh,
+# rng,
+# );
+#
+# closure.chain
+#
+# θ₀.layer_4.spectral_weights |> size
+
+# θ₀
+# θ₀ = 2*θ₀
+# θ₀.layer_6 ./= 20
+
+# Prepare training
+θ = T(1.0e-1) * device(θ₀);
+# θ = device(θ₀);
+# θ = 2 * device(θ₀);
+opt = Optimisers.setup(Adam(T(1.0e-2)), θ);
+callbackstate = Point2f[];
+it = rand(1:size(io_valid[1].u, 4), 50);
+validset = map(v -> v[:, :, :, it], io_valid[1]);
+
+# A-priori loss
+loss = createloss(mean_squared_error, closure);
+dataloader = createdataloader(io_train[1]; batchsize = 50, device);
+dataloaders = [dataloader]
+dataloader()
+
+# A-posteriori loss
+loss = IncompressibleNavierStokes.create_trajectory_loss(; setup, pressure_solver, closure);
+dataloaders = [
+ IncompressibleNavierStokes.createtrajectoryloader(data_train; device, nunroll = 20)
+ for _ = 1:4
+];
+loss(dataloaders[1](), device(θ₀))
+
+# Warm-up
+loss(dataloaders[1](), θ)
+@time loss(dataloaders[1](), θ);
+b = dataloaders[1]();
+first(gradient(θ -> loss(b, θ), θ));
+@time first(gradient(θ -> loss(b, θ), θ));
+GC.gc()
+CUDA.reclaim()
+
+map() do
+ i = 3
+ # h = 1000 * sqrt(eps(T))
+ h = cbrt(eps(T))
+ θ1 = copy(θ)
+ θ2 = copy(θ)
+ CUDA.@allowscalar θ1[i] -= h / 2
+ CUDA.@allowscalar θ2[i] += h / 2
+ b = dataloaders[1]()
+ @show loss(b, θ2) loss(b, θ1)
+ a = (loss(b, θ2) - loss(b, θ1)) / h
+ b = CUDA.@allowscalar first(gradient(θ -> loss(b, θ), θ))[i]
+ [a; b]
+end
+
+# Training
+# Note: The states `opt`, `θ`, and `callbackstate`
+# will not be overwritten until training is finished.
+# This allows for cancelling with "Control-C" should errors explode.
+(; opt, θ, callbackstate) = train(
+ dataloaders,
+ loss,
+ opt,
+ θ;
+ niter = 1000,
+ ncallback = 10,
+ callbackstate,
+ callback = create_callback(closure, device(validset)...; state = callbackstate),
+);
+GC.gc()
+CUDA.reclaim()
+
+Array(θ)
+
+# # Save trained parameters
+# jldsave("output/forced/theta_cnn.jld2"; theta = Array(θ))
+# jldsave("output/forced/theta_fno.jld2"; theta = Array(θ))
+
+# # Load trained parameters
+# θθ = load("output/theta_cnn.jld2")
+# θθ = load("output/theta_fno.jld2")
+# copyto!(θ, θθ["theta"])
+
+function relerr(u, uref, setup)
+ (; dimension, Ip) = setup.grid
+ D = dimension()
+ a, b = T(0), T(0)
+ for α = 1:D
+ a += sum(abs2, u[α][Ip] - uref[α][Ip])
+ b += sum(abs2, uref[α][Ip])
+ end
+ sqrt(a) / sqrt(b)
+end
+
+u, u₀ = nothing, nothing
+u = device.(data_test[1].u[1]);
+u₀ = device(data_test[1].u[1][1]);
+Δt = data_test[1].t[2] - data_test[1].t[1]
+tlims = extrema(data_test[1].t)
+length(u)
+length(data_test[1].t)
+
+state_nm, outputs = solve_unsteady(
+ setup,
+ u₀,
+ tlims;
+ Δt,
+ pressure_solver,
+ processors = (;
+ relerr = relerr_trajectory(u, setup),
+ log = timelogger(; nupdate = 1000),
+ ),
+)
+relerr_nm = outputs.relerr[]
+
+state_cnn, outputs = solve_unsteady(
+ (; setup..., closure_model = wrappedclosure(closure, θ, setup)),
+ u₀,
+ tlims;
+ Δt,
+ pressure_solver,
+ processors = (relerr = relerr_trajectory(u, setup), log = timelogger(; nupdate = 1)),
+)
+relerr_cnn = outputs.relerr[]
+
+relerr_nm
+relerr_cnn
+
+# dnm = relerr_nm
+# dcnn = relerr_cnn
+
+dnm
+dcnn
diff --git a/src/IncompressibleNavierStokes.jl b/src/IncompressibleNavierStokes.jl
index 71ce39e38..5e4322870 100644
--- a/src/IncompressibleNavierStokes.jl
+++ b/src/IncompressibleNavierStokes.jl
@@ -6,97 +6,44 @@ Energy-conserving solvers for the incompressible Navier-Stokes equations.
module IncompressibleNavierStokes
using Adapt
+using ChainRulesCore
using FFTW
using IterativeSolvers
+using KernelAbstractions
using LinearAlgebra
+using Makie
+using NNlib
using Printf
+using Random
using SparseArrays
+using StaticArrays
using Statistics
using WriteVTK: CollectionFile, paraview_collection, vtk_grid, vtk_save
-using Makie
-
-# Convenience notation
-const ⊗ = kron
-# Grid
-include("grid/dimension.jl")
-include("grid/grid.jl")
-include("grid/stretched_grid.jl")
-include("grid/cosine_grid.jl")
-include("grid/max_size.jl")
+# # Easily retrieve value from Val
+# (::Val{x})() where {x} = x
-# Models
-include("models/viscosity_models.jl")
-include("models/convection_models.jl")
-
-# Types
-include("force/force.jl")
-include("boundary_conditions/boundary_conditions.jl")
-include("operators/operators.jl")
+# General stuff
+include("boundary_conditions.jl")
+include("grid.jl")
include("setup.jl")
-
-# Boundary condtions
-include("boundary_conditions/bc_av3.jl")
-include("boundary_conditions/bc_av_stag3.jl")
-include("boundary_conditions/bc_diff3.jl")
-include("boundary_conditions/bc_diff_stag.jl")
-include("boundary_conditions/bc_diff_stag3.jl")
-include("boundary_conditions/bc_div2.jl")
-include("boundary_conditions/bc_general.jl")
-include("boundary_conditions/bc_general_stag.jl")
-include("boundary_conditions/bc_general_stag_diff.jl")
-include("boundary_conditions/bc_int2.jl")
-include("boundary_conditions/bc_int3.jl")
-include("boundary_conditions/bc_int_mixed2.jl")
-include("boundary_conditions/bc_int_mixed_stag2.jl")
-include("boundary_conditions/bc_int_mixed_stag3.jl")
-include("boundary_conditions/bc_vort3.jl")
-include("boundary_conditions/get_bc_vectors.jl")
-
-# Operators
-include("operators/operator_averaging.jl")
-include("operators/operator_convection_diffusion.jl")
-include("operators/operator_divergence.jl")
-include("operators/operator_interpolation.jl")
-include("operators/operator_postprocessing.jl")
-include("operators/operator_regularization.jl")
-include("operators/operator_turbulent_diffusion.jl")
-include("operators/operator_viscosity.jl")
-include("operators/operator_filter.jl")
-
-# Pressure solvers
-include("solvers/pressure/pressure_solvers.jl")
-include("solvers/pressure/pressure_poisson.jl")
-include("solvers/pressure/pressure_additional_solve.jl")
+include("pressure.jl")
+include("operators.jl")
# Time steppers
-include("momentum/momentumcache.jl")
include("time_steppers/methods.jl")
-include("time_steppers/tableaux.jl")
include("time_steppers/nstage.jl")
include("time_steppers/time_stepper_caches.jl")
include("time_steppers/step.jl")
include("time_steppers/isexplicit.jl")
include("time_steppers/lambda_max.jl")
+include("time_steppers/tableaux.jl")
# Preprocess
-include("preprocess/create_initial_conditions.jl")
+include("create_initial_conditions.jl")
# Processors
-include("processors/processors.jl")
-include("processors/real_time_plot.jl")
-include("processors/animator.jl")
-
-# Momentum equation
-include("momentum/compute_conservation.jl")
-include("momentum/check_symmetry.jl")
-include("momentum/convection_components.jl")
-include("momentum/convection.jl")
-include("momentum/diffusion.jl")
-include("momentum/momentum.jl")
-include("momentum/strain_tensor.jl")
-include("momentum/turbulent_K.jl")
-include("momentum/turbulent_viscosity.jl")
+include("processors.jl")
# Solvers
include("solvers/get_timestep.jl")
@@ -104,85 +51,45 @@ include("solvers/solve_steady_state.jl")
include("solvers/solve_unsteady.jl")
# Utils
-include("utils/filter_convection.jl")
+include("utils/plotgrid.jl")
+include("utils/save_vtk.jl")
include("utils/get_lims.jl")
include("utils/plotmat.jl")
+include("utils/spectral_stuff.jl")
-# Postprocess
-include("postprocess/get_velocity.jl")
-include("postprocess/get_vorticity.jl")
-include("postprocess/get_streamfunction.jl")
-include("postprocess/plot_force.jl")
-include("postprocess/plot_grid.jl")
-include("postprocess/plot_pressure.jl")
-include("postprocess/plot_velocity.jl")
-include("postprocess/plot_vorticity.jl")
-include("postprocess/plot_streamfunction.jl")
-include("postprocess/save_vtk.jl")
-
-# Force
-export SteadyBodyForce
-
-# Models
-export LaminarModel, MixingLengthModel, SmagorinskyModel, QRModel
-export NoRegConvectionModel, C2ConvectionModel, C4ConvectionModel, LerayConvectionModel
+# Boundary conditions
+export PeriodicBC, DirichletBC, SymmetricBC, PressureBC
# Processors
-export processor, step_logger, vtk_writer, field_saver
-export field_plotter, energy_history_plotter, energy_spectrum_plotter
+export processor, timelogger, vtk_writer, fieldsaver, realtimeplotter
+export fieldplot, energy_history_plot, energy_spectrum_plot
export animator
# Setup
-export Setup
-export operator_filter
+export Setup, temperature_equation
# 1D grids
export stretched_grid, cosine_grid
# Pressure solvers
-export DirectPressureSolver, CGPressureSolver, SpectralPressureSolver
-export pressure_poisson,
- pressure_poisson!, pressure_additional_solve, pressure_additional_solve!
+export default_psolver,
+ psolver_direct,
+ psolver_cg,
+ psolver_cg_matrix,
+ psolver_spectral,
+ psolver_spectral_lowmemory
-# Problems
+# Solvers
export solve_unsteady, solve_steady_state
-export momentum, momentum!
-export create_initial_conditions, random_field, get_bc_vectors, get_velocity
+# Field generation
+export create_initial_conditions, random_field
-export plot_force,
- plot_grid, plot_pressure, plot_streamfunction, plot_velocity, plot_vorticity, save_vtk
+# Utils
+export plotgrid, save_vtk
export plotmat
# ODE methods
-
-export AdamsBashforthCrankNicolsonMethod, OneLegMethod
-
-# Runge Kutta methods
-export ExplicitRungeKuttaMethod, ImplicitRungeKuttaMethod, runge_kutta_method
-
-# Explicit Methods
-export FE11, SSP22, SSP42, SSP33, SSP43, SSP104, rSSPs2, rSSPs3, Wray3, RK56, DOPRI6
-
-# Implicit Methods
-export BE11, SDIRK34, ISSPm2, ISSPs3
-
-# Half explicit methods
-export HEM3, HEM3BS, HEM5
-
-# Classical Methods
-export GL1, GL2, GL3, RIA1, RIA2, RIA3, RIIA1, RIIA2, RIIA3, LIIIA2, LIIIA3
-
-# Chebyshev methods
-export CHDIRK3, CHCONS3, CHC3, CHC5
-
-# Miscellaneous Methods
-export Mid22, MTE22, CN22, Heun33, RK33C2, RK33P2, RK44, RK44C2, RK44C23, RK44P2
-
-# DSRK Methods
-export DSso2, DSRK2, DSRK3
-
-# "Non-SSP" Methods of Wong & Spiteri
-export NSSP21, NSSP32, NSSP33, NSSP53
+export AdamsBashforthCrankNicolsonMethod, OneLegMethod, RKMethods
end
diff --git a/src/boundary_conditions.jl b/src/boundary_conditions.jl
new file mode 100644
index 000000000..890f46693
--- /dev/null
+++ b/src/boundary_conditions.jl
@@ -0,0 +1,510 @@
+abstract type AbstractBC end
+
+"""
+ PeriodicBC()
+
+Periodic boundary conditions. Must be periodic on both sides.
+"""
+struct PeriodicBC <: AbstractBC end
+
+"""
+ DirichletBC()
+
+No slip boundary conditions, where all velocity components are zero.
+
+ DirichletBC(u, dudt)
+
+Dirichlet boundary conditions for the velocity, where `u[1] = (x..., t) ->
+u1_BC` up to `u[d] = (x..., t) -> ud_BC`, where `d` is the dimension.
+
+To make the pressure the same order as velocity, also provide `dudt`.
+"""
+struct DirichletBC{F,G} <: AbstractBC
+ u::F
+ dudt::G
+end
+
+DirichletBC() = DirichletBC(nothing, nothing)
+DirichletBC(u) = DirichletBC(u, nothing)
+
+"""
+ SymmetricBC()
+
+Symmetric boundary conditions.
+The parallel velocity and pressure is the same at each side of the boundary.
+The normal velocity is zero.
+"""
+struct SymmetricBC <: AbstractBC end
+
+"""
+ PressureBC()
+
+Pressure boundary conditions.
+The pressure is prescribed on the boundary (usually an "outlet").
+The velocity has zero Neumann conditions.
+
+Note: Currently, the pressure is prescribed with the constant value of
+zero on the entire boundary.
+"""
+struct PressureBC <: AbstractBC end
+
+function ghost_a! end
+function ghost_b! end
+
+# Add opposite boundary ghost volume
+# Do everything in first function call for periodic
+function ghost_a!(::PeriodicBC, x)
+ Δx_a = x[2] - x[1]
+ Δx_b = x[end] - x[end-1]
+ pushfirst!(x, x[1] - Δx_b)
+ push!(x, x[end] + Δx_a)
+end
+ghost_b!(::PeriodicBC, x) = nothing
+
+# Add infinitely thin boundary volume
+ghost_a!(::DirichletBC, x) = pushfirst!(x, x[1])
+ghost_b!(::DirichletBC, x) = push!(x, x[end])
+
+# Duplicate boundary volume
+ghost_a!(::SymmetricBC, x) = pushfirst!(x, x[1] - (x[2] - x[1]))
+ghost_b!(::SymmetricBC, x) = push!(x, x[end] + (x[end] - x[end-1]))
+
+# Add infinitely thin boundary volume
+# On the left, we need to add two ghost volumes to have a normal component at
+# the left of the first ghost volume
+ghost_a!(::PressureBC, x) = pushfirst!(x, x[1], x[1])
+ghost_b!(::PressureBC, x) = push!(x, x[end])
+
+"""
+ offset_u(bc, isnormal, isright)
+
+Number of non-DOF velocity components at boundary.
+If `isnormal`, then the velocity is normal to the boundary, else parallel.
+If `isright`, it is at the end/right/rear/top boundary, otherwise beginning.
+"""
+function offset_u end
+
+"""
+ offset_p(bc)
+
+Number of non-DOF pressure components at boundary.
+"""
+function offset_p end
+
+offset_u(::PeriodicBC, isnormal, isright) = 1
+offset_p(::PeriodicBC, isright) = 1
+
+offset_u(::DirichletBC, isnormal, isright) = 1 + isnormal * isright
+offset_p(::DirichletBC, isright) = 1
+
+offset_u(::SymmetricBC, isnormal, isright) = 1 + isnormal * isright
+offset_p(::SymmetricBC, isright) = 1
+
+offset_u(::PressureBC, isnormal, isright) = 1 + !isnormal * !isright
+offset_p(::PressureBC, isright) = 1 + !isright
+
+function apply_bc_u! end
+function apply_bc_p! end
+function apply_bc_temp! end
+
+apply_bc_u(u, t, setup; kwargs...) = apply_bc_u!(copy.(u), t, setup; kwargs...)
+apply_bc_p(p, t, setup; kwargs...) = apply_bc_p!(copy(p), t, setup; kwargs...)
+apply_bc_temp(temp, t, setup; kwargs...) = apply_bc_temp!(copy(temp), t, setup; kwargs...)
+
+ChainRulesCore.rrule(::typeof(apply_bc_u), u, t, setup; kwargs...) = (
+ apply_bc_u(u, t, setup; kwargs...),
+ # With respect to (apply_bc_u, u, t, setup)
+ φbar -> (
+ NoTangent(),
+ # Important: identity operator should be part of `apply_bc_u_pullback`,
+ # but is actually implemented via the `copy` below instead.
+ apply_bc_u_pullback!(Tangent{typeof(u)}(copy.((φbar...,))...), t, setup; kwargs...),
+ NoTangent(),
+ NoTangent(),
+ ),
+)
+
+ChainRulesCore.rrule(::typeof(apply_bc_p), p, t, setup) = (
+ apply_bc_p(p, t, setup),
+ # With respect to (apply_bc_p, p, t, setup)
+ φbar -> (
+ NoTangent(),
+ apply_bc_p_pullback!(
+ # Important: identity operator should be part of `apply_bc_p_pullback`,
+ # but is actually implemented via the `copy` below instead.
+ copy(unthunk(φbar)),
+ t,
+ setup,
+ ),
+ NoTangent(),
+ NoTangent(),
+ ),
+)
+
+function apply_bc_u!(u, t, setup; kwargs...)
+ (; boundary_conditions) = setup
+ D = length(u)
+ for β = 1:D
+ apply_bc_u!(boundary_conditions[β][1], u, β, t, setup; isright = false, kwargs...)
+ apply_bc_u!(boundary_conditions[β][2], u, β, t, setup; isright = true, kwargs...)
+ end
+ u
+end
+
+function apply_bc_u_pullback!(φbar, t, setup; kwargs...)
+ (; grid, boundary_conditions) = setup
+ (; dimension) = grid
+ D = dimension()
+ for β = 1:D
+ apply_bc_u_pullback!(
+ boundary_conditions[β][1],
+ φbar,
+ β,
+ t,
+ setup;
+ isright = false,
+ kwargs...,
+ )
+ apply_bc_u_pullback!(
+ boundary_conditions[β][2],
+ φbar,
+ β,
+ t,
+ setup;
+ isright = true,
+ kwargs...,
+ )
+ end
+ φbar
+end
+
+function apply_bc_p!(p, t, setup; kwargs...)
+ (; boundary_conditions, grid) = setup
+ (; dimension) = grid
+ D = dimension()
+ for β = 1:D
+ apply_bc_p!(boundary_conditions[β][1], p, β, t, setup; isright = false)
+ apply_bc_p!(boundary_conditions[β][2], p, β, t, setup; isright = true)
+ end
+ p
+end
+
+function apply_bc_p_pullback!(φbar, t, setup; kwargs...)
+ (; grid, boundary_conditions) = setup
+ (; dimension) = grid
+ D = dimension()
+ for β = 1:D
+ apply_bc_p_pullback!(
+ boundary_conditions[β][1],
+ φbar,
+ β,
+ t,
+ setup;
+ isright = false,
+ kwargs...,
+ )
+ apply_bc_p_pullback!(
+ boundary_conditions[β][2],
+ φbar,
+ β,
+ t,
+ setup;
+ isright = true,
+ kwargs...,
+ )
+ end
+ φbar
+end
+
+function apply_bc_temp!(temp, t, setup; kwargs...)
+ (; temperature, grid) = setup
+ (; boundary_conditions) = temperature
+ (; dimension) = grid
+ D = dimension()
+ for β = 1:D
+ apply_bc_temp!(boundary_conditions[β][1], temp, β, t, setup; isright = false)
+ apply_bc_temp!(boundary_conditions[β][2], temp, β, t, setup; isright = true)
+ end
+ temp
+end
+
+function apply_bc_temp_pullback!(φbar, t, setup; kwargs...)
+ (; temperature, grid) = setup
+ (; boundary_conditions) = temperature
+ (; dimension) = grid
+ D = dimension()
+ for β = 1:D
+ apply_bc_temp_pullback!(
+ boundary_conditions[β][1],
+ φbar,
+ β,
+ t,
+ setup;
+ isright = false,
+ kwargs...,
+ )
+ apply_bc_temp_pullback!(
+ boundary_conditions[β][2],
+ φbar,
+ β,
+ t,
+ setup;
+ isright = true,
+ kwargs...,
+ )
+ end
+ φbar
+end
+
+function apply_bc_u!(::PeriodicBC, u, β, t, setup; isright, kwargs...)
+ (; grid, workgroupsize) = setup
+ (; dimension, N) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function _bc_a!(u, ::Val{α}, ::Val{β}) where {α,β}
+ I = @index(Global, Cartesian)
+ u[α][I] = u[α][I+(N[β]-2)*e(β)]
+ end
+ @kernel function _bc_b!(u, ::Val{α}, ::Val{β}) where {α,β}
+ I = @index(Global, Cartesian)
+ u[α][I+(N[β]-1)*e(β)] = u[α][I+e(β)]
+ end
+ ndrange = ntuple(γ -> γ == β ? 1 : N[γ], D)
+ for α = 1:D
+ if isright
+ _bc_b!(get_backend(u[1]), workgroupsize)(u, Val(α), Val(β); ndrange)
+ else
+ _bc_a!(get_backend(u[1]), workgroupsize)(u, Val(α), Val(β); ndrange)
+ end
+ end
+ u
+end
+
+function apply_bc_u_pullback!(::PeriodicBC, φbar, β, t, setup; isright, kwargs...)
+ (; grid, workgroupsize) = setup
+ (; dimension, N) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function adj_a!(φ, ::Val{α}, ::Val{β}) where {α,β}
+ I = @index(Global, Cartesian)
+ φ[α][I+(N[β]-2)*e(β)] += φ[α][I]
+ φ[α][I] = 0
+ end
+ @kernel function adj_b!(φ, ::Val{α}, ::Val{β}) where {α,β}
+ I = @index(Global, Cartesian)
+ φ[α][I+e(β)] += φ[α][I+(N[β]-1)*e(β)]
+ φ[α][I+(N[β]-1)*e(β)] = 0
+ end
+ ndrange = ntuple(γ -> γ == β ? 1 : N[γ], D)
+ for α = 1:D
+ if isright
+ adj_b!(get_backend(φbar[1]), workgroupsize)(φbar, Val(α), Val(β); ndrange)
+ else
+ adj_a!(get_backend(φbar[1]), workgroupsize)(φbar, Val(α), Val(β); ndrange)
+ end
+ end
+ φbar
+end
+
+function apply_bc_p!(::PeriodicBC, p, β, t, setup; isright, kwargs...)
+ (; grid, workgroupsize) = setup
+ (; dimension, N) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function _bc_a(p, ::Val{β}) where {β}
+ I = @index(Global, Cartesian)
+ p[I] = p[I+(N[β]-2)*e(β)]
+ end
+ @kernel function _bc_b(p, ::Val{β}) where {β}
+ I = @index(Global, Cartesian)
+ p[I+(N[β]-1)*e(β)] = p[I+e(β)]
+ end
+ ndrange = ntuple(γ -> γ == β ? 1 : N[γ], D)
+ if isright
+ _bc_b(get_backend(p), workgroupsize)(p, Val(β); ndrange)
+ else
+ _bc_a(get_backend(p), workgroupsize)(p, Val(β); ndrange)
+ end
+ p
+end
+
+function apply_bc_p_pullback!(::PeriodicBC, φbar, β, t, setup; isright, kwargs...)
+ (; grid, workgroupsize) = setup
+ (; dimension, N) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function adj_a!(φ, ::Val{β}) where {β}
+ I = @index(Global, Cartesian)
+ φ[I+(N[β]-2)*e(β)] += φ[I]
+ φ[I] = 0
+ end
+ @kernel function adj_b!(φ, ::Val{β}) where {β}
+ I = @index(Global, Cartesian)
+ φ[I+e(β)] += φ[I+(N[β]-1)*e(β)]
+ φ[I+(N[β]-1)*e(β)] = 0
+ end
+ ndrange = ntuple(γ -> γ == β ? 1 : N[γ], D)
+ if isright
+ adj_b!(get_backend(φbar), workgroupsize)(φbar, Val(β); ndrange)
+ else
+ adj_a!(get_backend(φbar), workgroupsize)(φbar, Val(β); ndrange)
+ end
+ φbar
+end
+
+apply_bc_temp!(bc::PeriodicBC, temp, β, t, setup; isright, kwargs...) =
+ apply_bc_p!(bc, temp, β, t, setup; isright, kwargs...)
+
+function apply_bc_u!(bc::DirichletBC, u, β, t, setup; isright, dudt = false, kwargs...)
+ (; dimension, x, xp, N) = setup.grid
+ D = dimension()
+ e = Offset{D}()
+ # isnothing(bc.u) && return
+ bcfunc = dudt ? bc.dudt : bc.u
+ for α = 1:D
+ I = if isright
+ CartesianIndices(
+ ntuple(γ -> γ == β ? α == β ? (N[γ]-1:N[γ]-1) : (N[γ]:N[γ]) : (1:N[γ]), D),
+ )
+ else
+ CartesianIndices(ntuple(γ -> γ == β ? (1:1) : (1:N[γ]), D))
+ end
+ xI = ntuple(
+ γ -> reshape(
+ γ == α ? x[γ][I.indices[α].+1] : xp[γ][I.indices[γ]],
+ ntuple(Returns(1), γ - 1)...,
+ :,
+ ntuple(Returns(1), D - γ)...,
+ ),
+ D,
+ )
+ if isnothing(bc.u)
+ u[α][I] .= 0
+ else
+ u[α][I] .= bcfunc.((Dimension(α),), xI..., t)
+ end
+ end
+ u
+end
+
+# apply_bc_u_pullback!(::DirichletBC, φbar, β, t, setup; isright, kwargs...) =
+# @not_implemented("DirichletBC pullback not yet implemented.")
+
+function apply_bc_p!(::DirichletBC, p, β, t, setup; isright, kwargs...)
+ (; dimension, N) = setup.grid
+ D = dimension()
+ e = Offset{D}()
+ if isright
+ I = CartesianIndices(ntuple(γ -> γ == β ? (N[γ]:N[γ]) : (1:N[γ]), D))
+ p[I] .= p[I.-e(β)]
+ else
+ I = CartesianIndices(ntuple(γ -> γ == β ? (1:1) : (1:N[γ]), D))
+ p[I] .= p[I.+e(β)]
+ end
+ p
+end
+
+# apply_bc_p_pullback!(::DirichletBC, φbar, β, t, setup; isright, kwargs...) =
+# @not_implemented("DirichletBC pullback not yet implemented.")
+
+function apply_bc_temp!(bc::DirichletBC, temp, β, t, setup; isright, kwargs...)
+ (; dimension, N) = setup.grid
+ D = dimension()
+ e = Offset{D}()
+ I = if isright
+ CartesianIndices(ntuple(γ -> γ == β ? (N[γ]:N[γ]) : (1:N[γ]), D))
+ else
+ CartesianIndices(ntuple(γ -> γ == β ? (1:1) : (1:N[γ]), D))
+ end
+ temp[I] .= bc.u
+ temp
+end
+
+function apply_bc_u!(::SymmetricBC, u, β, t, setup; isright, kwargs...)
+ (; dimension, N) = setup.grid
+ D = dimension()
+ e = Offset{D}()
+ for α = 1:D
+ if α != β
+ if isright
+ I = CartesianIndices(ntuple(γ -> γ == β ? (N[γ]:N[γ]) : (1:N[γ]), D))
+ u[α][I] .= u[α][I.-e(β)]
+ else
+ I = CartesianIndices(ntuple(γ -> γ == β ? (1:1) : (1:N[γ]), D))
+ u[α][I] .= u[α][I.+e(β)]
+ end
+ end
+ end
+ u
+end
+
+# apply_bc_u_pullback!(::SymmetricBC, φbar, β, t, setup; isright, kwargs...) =
+# @not_implemented("SymmetricBC pullback not yet implemented.")
+
+function apply_bc_p!(::SymmetricBC, p, β, t, setup; isright, kwargs...)
+ (; dimension, N) = setup.grid
+ D = dimension()
+ e = Offset{D}()
+ if isright
+ I = CartesianIndices(ntuple(γ -> γ == β ? (N[γ]:N[γ]) : (1:N[γ]), D))
+ p[I] .= p[I.-e(β)]
+ else
+ I = CartesianIndices(ntuple(γ -> γ == β ? (1:1) : (1:N[γ]), D))
+ p[I] .= p[I.+e(β)]
+ end
+ p
+end
+
+# apply_bc_p_pullback!(::SymmetricBC, φbar, β, t, setup; isright, kwargs...) =
+# @not_implemented("SymmetricBC pullback not yet implemented.")
+
+apply_bc_temp!(bc::SymmetricBC, temp, β, t, setup; isright, kwargs...) =
+ apply_bc_p!(bc, temp, β, t, setup; isright, kwargs...)
+
+function apply_bc_u!(bc::PressureBC, u, β, t, setup; isright, kwargs...)
+ (; grid, workgroupsize) = setup
+ (; dimension, N, Nu, Iu) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function _bc_a!(u, ::Val{α}, ::Val{β}, I0) where {α,β}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ u[α][I] = u[α][I+e(β)]
+ end
+ @kernel function _bc_b!(u, ::Val{α}, ::Val{β}, I0) where {α,β}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ u[α][I] = u[α][I-e(β)]
+ end
+ ndrange = (N[1:β-1]..., 1, N[β+1:end]...)
+ for α = 1:D
+ if isright
+ I0 = CartesianIndex(ntuple(γ -> γ == β ? N[β] : 1, D))
+ I0 -= oneunit(I0)
+ _bc_b!(get_backend(u[1]), workgroupsize)(u, Val(α), Val(β), I0; ndrange)
+ else
+ I0 = CartesianIndex(ntuple(γ -> γ == β && α != β ? 2 : 1, D))
+ I0 -= oneunit(I0)
+ _bc_a!(get_backend(u[1]), workgroupsize)(u, Val(α), Val(β), I0; ndrange)
+ end
+ end
+ u
+end
+
+# apply_bc_u_pullback!(::PressureBC, φbar, β, t, setup; isright, kwargs...) =
+# @not_implemented("PressureBC pullback not yet implemented.")
+
+function apply_bc_p!(bc::PressureBC, p, β, t, setup; isright, kwargs...)
+ (; dimension, N) = setup.grid
+ D = dimension()
+ I = if isright
+ CartesianIndices(ntuple(γ -> γ == β ? (N[γ]:N[γ]) : (1:N[γ]), D))
+ else
+ CartesianIndices(ntuple(γ -> γ == β ? (2:2) : (1:N[γ]), D))
+ end
+ p[I] .= 0
+ p
+end
+
+apply_bc_p_pullback!(::PressureBC, φbar, β, t, setup; isright, kwargs...) =
+ @not_implemented("PressureBC pullback not yet implemented.")
diff --git a/src/boundary_conditions/bc_av3.jl b/src/boundary_conditions/bc_av3.jl
deleted file mode 100644
index 4e7e3d9bc..000000000
--- a/src/boundary_conditions/bc_av3.jl
+++ /dev/null
@@ -1,93 +0,0 @@
-function bc_av3(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on boundary / grid lines
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0
- # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 6
- # normal situation, 2 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -3 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:3, 1:3] = I(3)
- Bb[(end-2):end, (end-2):end] = I(3)
- if bc1 == :dirichlet
- Bbc[1, 1] = 1 / 2
- Bbc[1, 5] = 1 / 2
- Bbc[2, 2] = 1 / 2
- Bbc[2, 4] = 1 / 2
- Bbc[3, 3] = 1 # Dirichlet uLe
- ybc1_1D[1] = 1
- ybc1_1D[2] = 1
- ybc1_1D[3] = 1
- elseif bc1 == :pressure
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, (end-5):(end-3)] = I(3)
- Bbc[(end-2):end, 4:6] = -I(3)
- Bbc[(end-2):end, (end-2):end] = I(3)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = 1 / 2
- Bbc[end, end-4] = 1 / 2
- Bbc[end-1, end-1] = 1 / 2
- Bbc[end-1, end-3] = 1 / 2
- Bbc[end-2, end-2] = 1
- ybc2_1D[end-2] = 1 # uRi
- ybc2_1D[end-1] = 1
- ybc2_1D[end] = 1
- elseif bc2 == :pressure
- error("not implemented")
- elseif bc2 == :periodic
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, (end-5):(end-3)] = I(3)
- Bbc[(end-2):end, 4:6] = -I(3)
- Bbc[(end-2):end, (end-2):end] = I(3)
- else
- error("not implemented")
- end
- elseif Nb == 1
- # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_av_stag3.jl b/src/boundary_conditions/bc_av_stag3.jl
deleted file mode 100644
index cb7091c57..000000000
--- a/src/boundary_conditions/bc_av_stag3.jl
+++ /dev/null
@@ -1,96 +0,0 @@
-function bc_av_stag3(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on staggered locations (pressure lines)
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0
- # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 6
- # normal situation, 2 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -3 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:3, 1:3] = I(3)
- Bb[(end-2):end, (end-2):end] = I(3)
-
- if bc1 == :dirichlet
- Bbc[1, 1] = -1
- Bbc[1, 6] = 1
- Bbc[2, 2] = 1 / 2 # Dirichlet uLo
- Bbc[2, 5] = 1 / 2
- Bbc[3, 3] = -1
- Bbc[3, 4] = 1
- ybc1_1D[1] = 0
- ybc1_1D[2] = 1
- ybc1_1D[3] = 0
- elseif bc1 == :symmetric
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, (end-5):(end-3)] = I(3)
- Bbc[(end-2):end, 4:6] = -I(3)
- Bbc[(end-2):end, (end-2):end] = I(3)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = -1
- Bbc[end, end-5] = 1
- Bbc[end-1, end-1] = 1 / 2
- Bbc[end-1, end-4] = 1 / 2
- Bbc[end-2, end-2] = -1
- Bbc[end-2, end-3] = 1
- ybc2_1D[end-2] = 0
- ybc2_1D[end-1] = 1
- ybc2_1D[end] = 0
- elseif bc2 == :symmetric
- error("not implemented")
- elseif bc2 == :periodic
- # actually redundant
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, (end-5):(end-3)] = I(3)
- Bbc[(end-2):end, 4:6] = -I(3)
- Bbc[(end-2):end, (end-2):end] = I(3)
- else
- error("not implemented")
- end
- elseif Nb == 1 # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_diff3.jl b/src/boundary_conditions/bc_diff3.jl
deleted file mode 100644
index 0f66cb28d..000000000
--- a/src/boundary_conditions/bc_diff3.jl
+++ /dev/null
@@ -1,93 +0,0 @@
-# total solution u is written as u = Bb*ub + Bin*uin
-# the boundary conditions can be written as Bbc*u = ybc
-# then u can be written entirely in terms of uin and ybc as:
-# u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
-# Btemp = Bb*(Bbc*Bb)^(-1)
-# Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-#
-# val1 and val2 can be scalars or vectors with either the value or the
-# derivative
-#
-# (ghost) points on boundary / grid lines
-function bc_diff3(Nt, Nin, Nb, bc1, bc2, h1, h2)
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0 # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 6
- # normal situation, 2 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -3 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:3, 1:3] = I(3)
- Bb[end-2:end, end-2:end] = I(3)
- if bc1 == :dirichlet
- Bbc[1, 1] = 1 / 2 # Dirichlet
- Bbc[1, 5] = 1 / 2
- Bbc[2, 2] = 1 / 2 # Dirichlet uLe
- Bbc[2, 4] = 1 / 2
- Bbc[3, 3] = 1 # Dirichlet uLe
- ybc1_1D[1] = 1
- ybc1_1D[2] = 1
- ybc1_1D[3] = 1
- elseif bc1 == :pressure
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, end-5:end-3] = I(3)
- Bbc[end-2:end, 4:6] = -I(3)
- Bbc[end-2:end, end-2:end] = I(3)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = 1 / 2
- Bbc[end, end-4] = 1 / 2
- Bbc[end-1, end-1] = 1 / 2
- Bbc[end-1, end-3] = 1 / 2
- Bbc[end-2, end-2] = 1
- ybc2_1D[end-2] = 1 # uRi
- ybc2_1D[end-1] = 1 # uRi
- ybc2_1D[end] = 1
- elseif bc2 == :pressure
- error("not implemented")
- elseif bc2 == :periodic
- # actually redundant
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, end-5:end-3] = I(3)
- Bbc[end-2:end, 4:6] = -I(3)
- Bbc[end-2:end, end-2:end] = I(3)
- else
- error("not implemented")
- end
-
- elseif Nb == 1
- # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_diff_stag.jl b/src/boundary_conditions/bc_diff_stag.jl
deleted file mode 100644
index a7260f2c6..000000000
--- a/src/boundary_conditions/bc_diff_stag.jl
+++ /dev/null
@@ -1,88 +0,0 @@
-function bc_diff_stag(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # Total solution u is written as u = Bb*ub + Bin*uin
- # The boundary conditions can be written as Bbc*u = ybc
- # Then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # Val1 and val2 can be scalars or vectors with either the value or the
- # Derivative
- # (ghost) points on staggered locations (pressure lines)
-
- T = typeof(h1)
-
- # Some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- # Boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 0
- # No boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- elseif Nb == 1
- # One boundary point (should not be unnecessary)
- elseif Nb == 2
- # Normal situation, 2 boundary points
-
- # Boundary matrices
- Bin = spdiagm(Nt, Nin, -1 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1, 1] = 1
- Bb[end, Nb] = 1
-
- if bc1 == :dirichlet
- # Zeroth order (standard mirror conditions)
- Bbc[1, 1] = 1 / 2
- Bbc[1, 2] = 1 / 2
- ybc1_1D[1] = 1 # ULo
- elseif bc1 == :symmetric
- Bbc[1, 1] = -1
- Bbc[1, 2] = 1
- ybc1_1D[1] = h1 # DuLo
- elseif bc1 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end-1] = 1
- Bbc[2, 2] = -1
- Bbc[2, end] = 1
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- # Zeroth order (standard mirror conditions)
- Bbc[end, end-1] = 1 / 2
- Bbc[end, end] = 1 / 2
- ybc2_1D[2] = 1 # UUp
- elseif bc2 == :symmetric
- Bbc[2, end-1] = -1
- Bbc[2, end] = 1
- ybc2_1D[2] = h2 # DuUp
- elseif bc2 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end-1] = 1
- Bbc[2, 2] = -1
- Bbc[2, end] = 1
- else
- error("not implemented")
- end
- end
-
- if Nb ∈ (1, 2)
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_diff_stag3.jl b/src/boundary_conditions/bc_diff_stag3.jl
deleted file mode 100644
index 69c3deea9..000000000
--- a/src/boundary_conditions/bc_diff_stag3.jl
+++ /dev/null
@@ -1,93 +0,0 @@
-function bc_diff_stag3(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # Total solution u is written as u = Bb*ub + Bin*uin
- # The boundary conditions can be written as Bbc*u = ybc
- # Then u can be written entirely in terms of uin and ybc as:
- # U = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # Val1 and val2 can be scalars or vectors with either the value or the
- # Derivative
- # (ghost) points on staggered locations (pressure lines)
-
- T = typeof(h1)
-
- # Some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- # Boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 0
- # No boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- elseif Nb == 6
- # Normal situation, 2 boundary points
- # Boundary matrices
- Bin = spdiagm(Nt, Nin, -3 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:3, 1:3] = I(3)
- Bb[(end-2):end, (end-2):end] = I(3)
-
- if bc1 == :dirichlet
- Bbc[1, 1] = 1 / 2 # Dirichlet
- Bbc[1, 6] = 1 / 2
- Bbc[2, 2] = 1 / 2 # Dirichlet uLo
- Bbc[2, 5] = 1 / 2
- Bbc[3, 3] = 1 / 2 # Dirichlet uLo
- Bbc[3, 4] = 1 / 2
- ybc1_1D[1] = 1
- ybc1_1D[2] = 1
- ybc1_1D[3] = 1
- elseif bc1 == :symmetric
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, (end-5):(end-3)] = I(3)
- Bbc[(end-2):end, 4:6] = -I(3)
- Bbc[(end-2):end, (end-2):end] = I(3)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = 1 / 2
- Bbc[end, end-5] = 1 / 2
- Bbc[end-1, end-1] = 1 / 2
- Bbc[end-1, end-4] = 1 / 2
- Bbc[end-2, end-2] = 1 / 2
- Bbc[end-2, end-3] = 1 / 2
- ybc2_1D[end-2] = 1
- ybc2_1D[end-1] = 1
- ybc2_1D[end] = 1
- elseif bc2 == :symmetric
- error("not implemented")
- elseif bc2 == :periodic
- # Actually redundant
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, (end-5):(end-3)] = I(3)
- Bbc[(end-2):end, 4:6] = -I(3)
- Bbc[(end-2):end, (end-2):end] = I(3)
- else
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- elseif Nb == 1
- # One boundary point
- error("not implemented")
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_div2.jl b/src/boundary_conditions/bc_div2.jl
deleted file mode 100644
index ef0816dee..000000000
--- a/src/boundary_conditions/bc_div2.jl
+++ /dev/null
@@ -1,87 +0,0 @@
-function bc_div2(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on boundary / grid lines
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0
- # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 4
- # normal situation, 4 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -2 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:2, 1:2] = I(2)
- Bb[end-1:end, end-1:end] = I(2)
- if bc1 == :dirichlet
- Bbc[1, 1] = 1 / 2
- Bbc[1, 3] = 1 / 2
- Bbc[2, 2] = 1 # Dirichlet uLe
- ybc1_1D[1] = 1
- ybc1_1D[2] = 1
- elseif bc1 == :pressure
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-3:end-2] = I(2)
- Bbc[end-1:end, 3:4] = -I(2)
- Bbc[end-1:end, end-1:end] = I(2)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = 1 / 2
- Bbc[end, end-2] = 1 / 2
- Bbc[end-1, end-1] = 1 # Dirichlet uRi
- ybc2_1D[end-1] = 1
- ybc2_1D[end] = 1
- elseif bc2 == :pressure
- error("not implemented")
- elseif bc2 == :periodic # actually redundant
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-3:end-2] = I(2)
- Bbc[end-1:end, 3:4] = -I(2)
- Bbc[end-1:end, end-1:end] = I(2)
- else
- error("not implemented")
- end
- else
- # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_general.jl b/src/boundary_conditions/bc_general.jl
deleted file mode 100644
index ae842697c..000000000
--- a/src/boundary_conditions/bc_general.jl
+++ /dev/null
@@ -1,119 +0,0 @@
-# bc_general(Nt, Nin, Nb, bc1, bc2, h1, h2)
-#
-# Total solution `u` is written as `u = Bb*ub + Bin*uin`
-#
-# The boundary conditions can be written as `Bbc*u = ybc`
-#
-# Then `u` can be written entirely in terms of `uin` and `ybc` as: `u =
-# (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc`, where `Btemp = Bb/(Bbc*Bb)`.
-#
-# `Bb`, `Bin` and `Bbc` depend on type of bc (Neumann/Dirichlet/periodic) `val1` and `val2`
-# can be scalars or vectors with either the value or the derivative (ghost) points on
-# boundary/grid lines
-function bc_general(Nt, Nin, Nb, bc1, bc2, h1, h2)
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- T = typeof(h1)
-
- # Boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 0
- # No boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- elseif Nb == 1
- # One boundary point
- Bb = spzeros(T, Nt, Nb)
- diagpos = -1
- if bc1 == :dirichlet
- Bbc[1, 1] = 1
- ybc1_1D[1] = 1 # uLe
- Bb[1, 1] = 1
- elseif bc1 == :pressure
- diagpos = 0
- elseif bc1 == :periodic
- diagpos = 0
- Bbc[1, 1] = -1
- Bbc[1, end] = 1
- Bb[end, 1] = 1
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[Nb, end] = 1
- ybc2_1D[1] = 1 # uRi
- Bb[end, Nb] = 1
- elseif bc2 == :pressure
-
- elseif bc2 == :periodic # Actually redundant
- diagpos = 0
- Bbc[1, 1] = -1
- Bbc[1, end] = 1
- Bb[end, 1] = 1
- else
- error("not implemented")
- end
-
- # Boundary matrices
- Bin = spdiagm(Nt, Nin, diagpos => ones(T, Nin))
- elseif Nb == 2
- # Normal situation, 2 boundary points
- # Boundary matrices
- Bin = spdiagm(Nt, Nin, -1 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1, 1] = 1
- Bb[end, Nb] = 1
-
- if bc1 == :dirichlet
- Bbc[1, 1] = 1
- ybc1_1D[1] = 1 # uLe
- elseif bc1 == :pressure
- Bbc[1, 1] = -1
- Bbc[1, 2] = 1
- ybc1_1D[1] = 2 * h1 # DuLe
- elseif bc1 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end-1] = 1
- Bbc[Nb, 2] = -1
- Bbc[Nb, end] = 1
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[Nb, end] = 1
- ybc2_1D[2] = 1 # uRi
- elseif bc2 == :pressure
- Bbc[Nb, end-1] = -1
- Bbc[Nb, end] = 1
- ybc2_1D[2] = 2 * h2 # duRi
- elseif bc2 == :periodic # Actually redundant
- Bbc[1, 1] = -1
- Bbc[1, end-1] = 1
- Bbc[Nb, 2] = -1
- Bbc[Nb, end] = 1
- else
- error("not implemented")
- end
- else
- error("Nb must be 0, 1, or 2")
- end
-
- if Nb ∈ (1, 2)
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_general_stag.jl b/src/boundary_conditions/bc_general_stag.jl
deleted file mode 100644
index b2d9c19c8..000000000
--- a/src/boundary_conditions/bc_general_stag.jl
+++ /dev/null
@@ -1,124 +0,0 @@
-function bc_general_stag(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # Total solution u is written as u = Bb*ub + Bin*uin
- # The boundary conditions can be written as Bbc*u = ybc
- # Then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1]
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
- # Val1 and val2 can be scalars or vectors with either the value or the
- # derivative
- # (ghost) points on staggered locations (pressure lines)
-
- T = typeof(h1)
-
- # Some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- # Boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 0
- # No boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- elseif Nb == 1
- # One boundary point
- Bb = spzeros(T, Nt, Nb)
-
- diagpos = -1
- if bc1 == :dirichlet
- Bbc[1, 1] = 1 / 2
- Bbc[1, 2] = 1 / 2
- ybc1_1D[1] = 1 # ULe
- Bb[1, 1] = 1
- elseif bc1 == :symmetric
- Bbc[1, 1] = -1
- Bbc[1, 2] = 1
- ybc1_1D[1] = h1 # DuLo
- elseif bc1 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end] = 1
- Bb[1, 1] = 1
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[Nb, end-1] = 1 / 2
- Bbc[Nb, end] = 1 / 2
- ybc2_1D[1] = 1 # URi
- Bb[end, Nb] = 1
- elseif bc2 == :symmetric
- Bbc[Nb, end-1] = -1
- Bbc[Nb, end] = 1
- ybc2_1D[1] = h2 # DuUp
- elseif bc2 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end] = 1
- Bb[1, 1] = 1
- else
- error("not implemented")
- end
-
- # Boundary matrices
- Bin = spdiagm(Nt, Nin, diagpos => ones(T, Nin))
- elseif Nb == 2
- # Normal situation, 2 boundary points
- # Boundary matrices
- Bin = spdiagm(Nt, Nin, -1 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1, 1] = 1
- Bb[end, Nb] = 1
-
- if bc1 == :dirichlet
- Bbc[1, 1] = 1 / 2
- Bbc[1, 2] = 1 / 2
- ybc1_1D[1] = 1 # ULo
- elseif bc1 == :symmetric
- Bbc[1, 1] = -1
- Bbc[1, 2] = 1
- ybc1_1D[1] = h1 # DuLo
- elseif bc1 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end-1] = 1
- Bbc[2, 2] = -1
- Bbc[2, end] = 1
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end-1] = 1 / 2
- Bbc[end, end] = 1 / 2
- ybc2_1D[2] = 1 # UUp
- elseif bc2 == :symmetric
- Bbc[2, end-1] = -1
- Bbc[2, end] = 1
- ybc2_1D[2] = h2 # DuUp
- elseif bc2 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end-1] = 1
- Bbc[2, 2] = -1
- Bbc[2, end] = 1
- else
- error("not implemented")
- end
- else
- error("Nb must be 0, 1, or 2")
- end
-
- if Nb ∈ (1, 2)
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_general_stag_diff.jl b/src/boundary_conditions/bc_general_stag_diff.jl
deleted file mode 100644
index a488bda8c..000000000
--- a/src/boundary_conditions/bc_general_stag_diff.jl
+++ /dev/null
@@ -1,88 +0,0 @@
-function bc_general_stag_diff(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # Total solution u is written as u = Bb*ub + Bin*uin
- # The boundary conditions can be written as Bbc*u = ybc
- # Then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # Val1 and val2 can be scalars or vectors with either the value or the
- # Derivative
- # (ghost) points on staggered locations (pressure lines)
-
- T = typeof(h1)
-
- # Some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- # Boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 0
- # No boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- elseif Nb == 1
- # One boundary point (should not be unnecessary)
- elseif Nb == 2
- # Normal situation, 2 boundary points
-
- # Boundary matrices
- Bin = spdiagm(Nt, Nin, -1 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1, 1] = 1
- Bb[end, Nb] = 1
-
- if bc1 == :dirichlet
- Bbc[1, 1] = 3 / 8
- Bbc[1, 2] = 3 / 4
- Bbc[1, 3] = -1 / 8
- ybc1_1D[1] = 1 # ULo
- elseif bc1 == :symmetric
- Bbc[1, 1] = -1
- Bbc[1, 2] = 1
- ybc1_1D[1] = h1 # DuLo
- elseif bc1 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end-1] = 1
- Bbc[2, 2] = -1
- Bbc[2, end] = 1
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end-2] = -1 / 8
- Bbc[end, end-1] = 3 / 4
- Bbc[end, end] = 3 / 8
- ybc2_1D[2] = 1 # UUp
- elseif bc2 == :symmetric
- Bbc[2, end-1] = -1
- Bbc[2, end] = 1
- ybc2_1D[2] = h2 # DuUp
- elseif bc2 == :periodic
- Bbc[1, 1] = -1
- Bbc[1, end-1] = 1
- Bbc[2, 2] = -1
- Bbc[2, end] = 1
- else
- error("not implemented")
- end
- end
-
- if Nb ∈ (1, 2)
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_int2.jl b/src/boundary_conditions/bc_int2.jl
deleted file mode 100644
index 2c8c8a5f6..000000000
--- a/src/boundary_conditions/bc_int2.jl
+++ /dev/null
@@ -1,86 +0,0 @@
-function bc_int2(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on boundary / grid lines
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0 # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 4
- # normal situation, 4 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -2 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:2, 1:2] = I(2)
- Bb[end-1:end, end-1:end] = I(2)
- if bc1 == :dirichlet
- Bbc[1, 1] = 1 / 2
- Bbc[1, 3] = 1 / 2
- Bbc[2, 2] = 1 # Dirichlet uLe
- ybc1_1D[1] = 1
- ybc1_1D[2] = 1
- elseif bc1 == :pressure
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-3:end-2] = I(2)
- Bbc[end-1:end, 3:4] = -I(2)
- Bbc[end-1:end, end-1:end] = I(2)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = 1 / 2
- Bbc[end, end-2] = 1 / 2
- Bbc[end-1, end-1] = 1
- ybc2_1D[end-1] = 1 # uRi
- ybc2_1D[end] = 1
- elseif bc2 == :pressure
- error("not implemented")
- elseif bc2 == :periodic # actually redundant
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-3:end-2] = I(2)
- Bbc[end-1:end, 3:4] = -I(2)
- Bbc[end-1:end, end-1:end] = I(2)
- else
- error("not implemented")
- end
-
- elseif Nb == 1 # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_int3.jl b/src/boundary_conditions/bc_int3.jl
deleted file mode 100644
index cf8421650..000000000
--- a/src/boundary_conditions/bc_int3.jl
+++ /dev/null
@@ -1,93 +0,0 @@
-function bc_int3(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on boundary / grid lines
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0
- # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 6
- # normal situation, 2 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -3 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:3, 1:3] = I(3)
- Bb[end-2:end, end-2:end] = I(3)
- if bc1 == :dirichlet
- Bbc[1, 1] = 1 / 2 # skew-symm
- Bbc[1, 5] = 1 / 2
- Bbc[2, 2] = 1 / 2
- Bbc[2, 4] = 1 / 2
- Bbc[3, 3] = 1 # Dirichlet uLe
- ybc1_1D[1] = 1
- ybc1_1D[2] = 1
- ybc1_1D[3] = 1
- elseif bc1 == :pressure
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, end-5:end-3] = I(3)
- Bbc[end-2:end, 4:6] = -I(3)
- Bbc[end-2:end, end-2:end] = I(3)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = 1 / 2
- Bbc[end, end-4] = 1 / 2
- Bbc[end-1, end-1] = 1 / 2
- Bbc[end-1, end-3] = 1 / 2
- Bbc[end-2, end-2] = 1
- ybc2_1D[end-2] = 1 # uRi
- ybc2_1D[end-1] = 1
- ybc2_1D[end] = 1
- elseif bc2 == :pressure
- error("not implemented")
- elseif bc2 == :periodic # actually redundant
- Bbc[1:3, 1:3] = -I(3)
- Bbc[1:3, end-5:end-3] = I(3)
- Bbc[end-2:end, 4:6] = -I(3)
- Bbc[end-2:end, end-2:end] = I(3)
- else
- error("not implemented")
- end
- elseif Nb == 1
- # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_int_mixed2.jl b/src/boundary_conditions/bc_int_mixed2.jl
deleted file mode 100644
index 917458ab2..000000000
--- a/src/boundary_conditions/bc_int_mixed2.jl
+++ /dev/null
@@ -1,74 +0,0 @@
-function bc_int_mixed2(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on boundary / grid lines
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0 # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if bc1 == :dirichlet && bc2 == :dirichlet
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -2 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:2, 1:2] = I(2)
- Bb[end-1:end, end-1:end] = I(2)
-
- Bbc[1, 1] = -1
- Bbc[1, 3] = 1
- Bbc[2, 2] = 1 # Dirichlet uLe
- ybc1_1D[1] = 0
- ybc1_1D[2] = 1
-
- Bbc[end, end] = -1
- Bbc[end, end-2] = 1
- Bbc[end-1, end-1] = 1
- ybc2_1D[end-1] = 1 # uRi
- ybc2_1D[end] = 0
- elseif bc1 == :periodic && bc2 == :periodic
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -1 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1, 1] = 1
- Bb[end-1:end, end-1:end] = I(2)
- Bbc[1, 1] = -1
- Bbc[1, end-2] = 1
- Bbc[end-1, 2] = -1
- Bbc[end-1, end-1] = 1
- Bbc[end, 3] = -1
- Bbc[end, end] = 1
- else
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_int_mixed_stag2.jl b/src/boundary_conditions/bc_int_mixed_stag2.jl
deleted file mode 100644
index c14843786..000000000
--- a/src/boundary_conditions/bc_int_mixed_stag2.jl
+++ /dev/null
@@ -1,90 +0,0 @@
-function bc_int_mixed_stag2(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on staggered locations (pressure lines)
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0
- # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 4
- # normal situation, 2 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -2 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:2, 1:2] = I(2)
- Bb[end-1:end, end-1:end] = I(2)
-
- if bc1 == :dirichlet
- Bbc[1, 1] = -1 # Neumann type for skew-symmetry
- Bbc[1, 4] = 1
- Bbc[2, 2] = -1
- Bbc[2, 3] = 1
- ybc1_1D[1] = 0
- ybc1_1D[2] = 0
- elseif bc1 == :symmetric
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-3:end-2] = I(2)
- Bbc[end-1:end, 3:4] = -I(2)
- Bbc[end-1:end, end-1:end] = I(2)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = -1 # Neumann type for skew-symmetry
- Bbc[end, end-3] = 1
- Bbc[end-1, end-1] = -1
- Bbc[end-1, end-2] = 1
- ybc2_1D[end-1] = 0
- ybc2_1D[end] = 0
- elseif bc2 == :symmetric
- error("not implemented")
- elseif bc2 == :periodic # actually redundant
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-3:end-2] = I(2)
- Bbc[end-1:end, 3:4] = -I(2)
- Bbc[end-1:end, end-1:end] = I(2)
- else
- error("not implemented")
- end
- elseif Nb == 1
- # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_int_mixed_stag3.jl b/src/boundary_conditions/bc_int_mixed_stag3.jl
deleted file mode 100644
index eb0c448d7..000000000
--- a/src/boundary_conditions/bc_int_mixed_stag3.jl
+++ /dev/null
@@ -1,89 +0,0 @@
-function bc_int_mixed_stag3(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on staggered locations (pressure lines)
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0 # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 4
- # normal situation, 2 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -2 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:2, 1:2] = I(2)
- Bb[end-1:end, end-1:end] = I(2)
-
- if bc1 == :dirichlet
- Bbc[1, 1] = -1 # Neumann type for skew-symmetry
- Bbc[1, 4] = 1
- Bbc[2, 2] = -1
- Bbc[2, 3] = 1
- ybc1_1D[1] = 0
- ybc1_1D[2] = 0
- elseif bc1 == :symmetric
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-3:end-2] = I(2)
- Bbc[end-1:end, 3:4] = -I(2)
- Bbc[end-1:end, end-1:end] = I(2)
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- Bbc[end, end] = -1 # Neumann type for skew-symmetry
- Bbc[end, end-3] = 1
- Bbc[end-1, end-1] = -1
- Bbc[end-1, end-2] = 1
- ybc2_1D[end-1] = 0
- ybc2_1D[end] = 0
- elseif bc2 == :symmetric
- error("not implemented")
- elseif bc2 == :periodic # actually redundant
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-3:end-2] = I(2)
- Bbc[end-1:end, 3:4] = -I(2)
- Bbc[end-1:end, end-1:end] = I(2)
- else
- error("not implemented")
- end
- elseif Nb == 1
- # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/bc_vectors.jl b/src/boundary_conditions/bc_vectors.jl
deleted file mode 100644
index aec7bcfc7..000000000
--- a/src/boundary_conditions/bc_vectors.jl
+++ /dev/null
@@ -1,87 +0,0 @@
-Base.@kwdef struct BCVectors{T}
- yM::Vector{T} = T[]
- y_p::Vector{T} = T[]
-
- yAu_ux::Vector{T} = T[]
- yAu_uy::Vector{T} = T[]
- yAu_uz::Vector{T} = T[]
- yAv_vx::Vector{T} = T[]
- yAv_vy::Vector{T} = T[]
- yAv_vz::Vector{T} = T[]
- yAw_wx::Vector{T} = T[]
- yAw_wy::Vector{T} = T[]
- yAw_wz::Vector{T} = T[]
-
- yDiff::Vector{T} = T[]
-
- yIu_ux::Vector{T} = T[]
- yIv_uy::Vector{T} = T[]
- yIw_uz::Vector{T} = T[]
- yIu_vx::Vector{T} = T[]
- yIv_vy::Vector{T} = T[]
- yIw_vz::Vector{T} = T[]
- yIu_wx::Vector{T} = T[]
- yIv_wy::Vector{T} = T[]
- yIw_wz::Vector{T} = T[]
-
- ySu_ux::Vector{T} = T[]
- ySu_uy::Vector{T} = T[]
- ySu_uz::Vector{T} = T[]
- ySu_vx::Vector{T} = T[]
- ySu_wx::Vector{T} = T[]
-
- ySv_vx::Vector{T} = T[]
- ySv_vy::Vector{T} = T[]
- ySv_vz::Vector{T} = T[]
- ySv_uy::Vector{T} = T[]
- ySv_wy::Vector{T} = T[]
-
- ySw_wx::Vector{T} = T[]
- ySw_wy::Vector{T} = T[]
- ySw_wz::Vector{T} = T[]
- ySw_uz::Vector{T} = T[]
- ySw_vz::Vector{T} = T[]
-
- ydM::Vector{T} = T[]
-
- yAν_ux::Vector{T} = T[]
- yAν_uy::Vector{T} = T[]
- yAν_uz::Vector{T} = T[]
- yAν_vx::Vector{T} = T[]
- yAν_vy::Vector{T} = T[]
- yAν_vz::Vector{T} = T[]
- yAν_wx::Vector{T} = T[]
- yAν_wy::Vector{T} = T[]
- yAν_wz::Vector{T} = T[]
-
- yCux_k::Vector{T} = T[]
- yCuy_k::Vector{T} = T[]
- yCuz_k::Vector{T} = T[]
- yCvx_k::Vector{T} = T[]
- yCvy_k::Vector{T} = T[]
- yCvz_k::Vector{T} = T[]
- yCwx_k::Vector{T} = T[]
- yCwy_k::Vector{T} = T[]
- yCwz_k::Vector{T} = T[]
-
- yAuy_k::Vector{T} = T[]
- yAuz_k::Vector{T} = T[]
- yAvx_k::Vector{T} = T[]
- yAvz_k::Vector{T} = T[]
- yAwx_k::Vector{T} = T[]
- yAwy_k::Vector{T} = T[]
-
- yDiffu_f::Vector{T} = T[]
- yDiffv_f::Vector{T} = T[]
- yDiffw_f::Vector{T} = T[]
-
- yIu_ux3::Vector{T} = T[]
- yIv_uy3::Vector{T} = T[]
- yIu_vx3::Vector{T} = T[]
- yIv_vy3::Vector{T} = T[]
-
- yAu_ux3::Vector{T} = T[]
- yAu_uy3::Vector{T} = T[]
- yAv_vx3::Vector{T} = T[]
- yAv_vy3::Vector{T} = T[]
-end
diff --git a/src/boundary_conditions/bc_vort3.jl b/src/boundary_conditions/bc_vort3.jl
deleted file mode 100644
index 4fdd9edca..000000000
--- a/src/boundary_conditions/bc_vort3.jl
+++ /dev/null
@@ -1,83 +0,0 @@
-function bc_vort3(Nt, Nin, Nb, bc1, bc2, h1, h2)
- # total solution u is written as u = Bb*ub + Bin*uin
- # the boundary conditions can be written as Bbc*u = ybc
- # then u can be written entirely in terms of uin and ybc as:
- # u = (Bin-Btemp*Bbc*Bin)*uin + Btemp*ybc, where
- # Btemp = Bb*(Bbc*Bb)^(-1)
- # Bb, Bin and Bbc depend on type of bc (Neumann/Dirichlet/periodic)
-
- # val1 and val2 can be scalars or vectors with either the value or the
- # derivative
-
- # (ghost) points on boundary / grid lines
-
- # to calculate vorticity; two ghost points at lower/left boundary, one
- # ghost point upper/right boundary
-
- T = typeof(h1)
-
- # some input checking:
- if Nt != Nin + Nb
- error("Number of inner points plus boundary points is not equal to total points")
- end
-
- if Nb == 0
- # no boundary points, so simply diagonal matrix without boundary contribution
- B1D = I(Nt)
- Btemp = spzeros(T, Nt, 2)
- ybc1 = zeros(T, 2)
- ybc2 = zeros(T, 2)
- else
- # boundary conditions
- Bbc = spzeros(T, Nb, Nt)
- ybc1_1D = zeros(T, Nb)
- ybc2_1D = zeros(T, Nb)
-
- if Nb == 3
- # normal situation, 4 boundary points
-
- # boundary matrices
- Bin = spdiagm(Nt, Nin, -2 => ones(T, Nin))
- Bb = spzeros(T, Nt, Nb)
- Bb[1:2, 1:2] = I(2)
- Bb[end, end] = 1
- if bc1 == :dirichlet
- error("not implemented")
- elseif bc1 == :pressure
- error("not implemented")
- elseif bc1 == :periodic
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-2:end-1] = I(2)
- Bbc[end, 3] = -1
- Bbc[end, end] = 1
- else
- error("not implemented")
- end
-
- if bc2 == :dirichlet
- error("not implemented")
- elseif bc2 == :pressure
- error("not implemented")
- elseif bc2 == :periodic
- # actually redundant
- Bbc[1:2, 1:2] = -I(2)
- Bbc[1:2, end-2:end-1] = I(2)
- Bbc[end, 3] = -1
- Bbc[end, end] = 1
- else
- error("not implemented")
- end
- else
- # one boundary point
- error("not implemented")
- end
-
- ybc1 = ybc1_1D
- ybc2 = ybc2_1D
-
- Btemp = Bb / (Bbc * Bb)
- B1D = Bin - Btemp * Bbc * Bin
- end
-
- (; B1D, Btemp, ybc1, ybc2)
-end
diff --git a/src/boundary_conditions/boundary_conditions.jl b/src/boundary_conditions/boundary_conditions.jl
deleted file mode 100644
index 266de1025..000000000
--- a/src/boundary_conditions/boundary_conditions.jl
+++ /dev/null
@@ -1,95 +0,0 @@
-"""
- BoundaryConditions(u_bc, v_bc; T = Float64, bc_unsteady, bc_type, kwargs...)
-
-Create discrete boundary condtions.
-
-Values should either be scalars or vectors. All values `(u, v, p, k, e)` are
-defined at (x, y) locations, i.e. the corners of pressure volumes, so they
-cover the entire domain, including corners.
-"""
-function BoundaryConditions(
- u_bc,
- v_bc;
- T = Float64,
- bc_type,
- dudt_bc = nothing,
- dvdt_bc = nothing,
- bc_unsteady = !isnothing(dudt_bc),
- kwargs...,
-)
- bc_type.u.x[1] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for u-left")
- bc_type.u.x[2] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for u-right")
- bc_type.u.y[1] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for u-low")
- bc_type.u.y[2] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for u-up")
-
- bc_type.v.x[1] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for v-left")
- bc_type.v.x[2] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for v-right")
- bc_type.v.y[1] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for v-low")
- bc_type.v.y[2] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for v-up")
-
- # Pressure (for boundaries marked with `:pressure`)
- p∞ = zero(T)
- p_bc = (; x = (p∞, p∞), y = (p∞, p∞))
-
- (; bc_unsteady, bc_type..., u_bc, v_bc, dudt_bc, dvdt_bc, p_bc, kwargs...)
-end
-
-"""
- BoundaryConditions(u_bc, v_bc, w_bc; T = Float64, bc_unsteady, bc_type, kwargs...)
-
-Create discrete boundary condtions.
-
-Values should either be scalars or vectors. All values `(u, v, p, k, e)` are
-defined at (x, y, z) locations, i.e. the corners of pressure volumes, so they
-cover the entire domain, including corners.
-"""
-function BoundaryConditions(
- u_bc,
- v_bc,
- w_bc;
- T = Float64,
- bc_type,
- dudt_bc = nothing,
- dvdt_bc = nothing,
- dwdt_bc = nothing,
- bc_unsteady = !isnothing(dudt_bc),
- kwargs...,
-)
- bc_type.u.x[1] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for u-left")
- bc_type.u.x[2] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for u-right")
- bc_type.u.y[1] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for u-low")
- bc_type.u.y[2] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for u-up")
- bc_type.u.z[1] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for u-back")
- bc_type.u.z[2] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for u-front")
-
- bc_type.v.x[1] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for v-left")
- bc_type.v.x[2] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for v-right")
- bc_type.v.y[1] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for v-low")
- bc_type.v.y[2] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for v-up")
- bc_type.v.z[1] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for v-back")
- bc_type.v.z[2] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for v-front")
-
- bc_type.w.x[1] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for w-left")
- bc_type.w.x[2] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for w-right")
- bc_type.w.y[1] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for w-low")
- bc_type.w.y[2] ∈ (:dirichlet, :periodic, :symmetric) || error("Wrong BC for w-up")
- bc_type.w.z[1] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for w-back")
- bc_type.w.z[2] ∈ (:dirichlet, :periodic, :pressure) || error("Wrong BC for w-front")
-
- # Pressure (for boundaries marked with `:pressure`)
- p∞ = zero(T)
- p_bc = (; x = (p∞, p∞), y = (p∞, p∞), z = (p∞, p∞))
-
- (;
- bc_unsteady,
- bc_type...,
- u_bc,
- v_bc,
- w_bc,
- dudt_bc,
- dvdt_bc,
- dwdt_bc,
- p_bc,
- kwargs...,
- )
-end
diff --git a/src/boundary_conditions/get_bc_vectors.jl b/src/boundary_conditions/get_bc_vectors.jl
deleted file mode 100644
index e718a7841..000000000
--- a/src/boundary_conditions/get_bc_vectors.jl
+++ /dev/null
@@ -1,1024 +0,0 @@
-"""
- get_bc_vectors(setup, t)
-
-Get boundary condition vectors.
-"""
-function get_bc_vectors end
-
-get_bc_vectors(setup, t) = get_bc_vectors(setup.grid.dimension, setup, t)
-
-# 2D version
-function get_bc_vectors(::Dimension{2}, setup, t)
- (; grid, operators, boundary_conditions, viscosity_model) = setup
-
- (; Nux_in, Nvy_in, Np, Npx, Npy) = grid
- (; xin, yin, x, y, hx, hy, xp, yp) = grid
- (; Ω, indu, indv) = grid
- (; order4, α) = grid
-
- (; Dux, Duy, Dvx, Dvy) = operators
- (; Au_ux_bc, Au_uy_bc, Av_vx_bc, Av_vy_bc) = operators
- (; Su_ux_bc, Su_uy_bc, Sv_vx_bc, Sv_vy_bc) = operators
- (; Iu_ux_bc, Iv_uy_bc_lr, Iv_uy_bc_lu) = operators
- (; Iu_vx_bc_lr, Iu_vx_bc_lu, Iv_vy_bc) = operators
- (; Mx_bc, My_bc) = operators
- (; Su_vx_bc_lr, Su_vx_bc_lu, Sv_uy_bc_lr, Sv_uy_bc_lu) = operators
-
- (; u_bc, v_bc, dudt_bc, dvdt_bc) = boundary_conditions
- (; p_bc, bc_unsteady) = boundary_conditions
-
- (; Re) = viscosity_model
-
- if order4
- (; Au_ux_bc3, Au_uy_bc3, Av_vx_bc3, Av_vy_bc3) = operators
- (; Iu_ux_bc3, Iv_uy_bc_lu3, Iv_uy_bc_lr3) = operators
- (; Iu_vx_bc_lu3, Iu_vx_bc_lr3, Iv_vy_bc3) = operators
- (; Su_ux_bc3, Su_uy_bc3, Sv_vx_bc3, Sv_vy_bc3) = operators
- (; Diffux_div, Diffuy_div, Diffvx_div, Diffvy_div) = operators
- (; Mx_bc3, My_bc3) = operators
- end
-
- T = typeof(t)
-
- # TODO: Split function into allocating part (constructor?) and mutating `update!`
-
- ## Get BC values
- uLo = u_bc.(x, y[1], t)
- uUp = u_bc.(x, y[end], t)
-
- uLo_i = u_bc.(xin, y[1], t)
- uUp_i = u_bc.(xin, y[end], t)
- uLe_i = u_bc.(x[1], yp, t)
- uRi_i = u_bc.(x[end], yp, t)
-
- vLe = v_bc.(x[1], y, t)
- vRi = v_bc.(x[end], y, t)
-
- vLo_i = v_bc.(xp, y[1], t)
- vUp_i = v_bc.(xp, y[end], t)
- vLe_i = v_bc.(x[1], yin, t)
- vRi_i = v_bc.(x[end], yin, t)
-
- if bc_unsteady
- dudtLe_i = dudt_bc.(x[1], yp, t)
- dudtRi_i = dudt_bc.(x[end], yp, t)
- dvdtLo_i = dvdt_bc.(xp, y[1], t)
- dvdtUp_i = dvdt_bc.(xp, y[end], t)
- end
-
- ## Boundary conditions for divergence
-
- # Mx
- # ybc = uLe_i ⊗ Mx_bc.ybc1 + uRi_i ⊗ Mx_bc.ybc2
- ybc = vec(Mx_bc.ybc1 * uLe_i') + vec(Mx_bc.ybc2 * uRi_i')
- yMx = Mx_bc.Bbc * ybc
- if order4
- # ybc3 = uLe_i ⊗ Mx_bc3.ybc1 + uRi_i ⊗ Mx_bc3.ybc2
- ybc3 = vec(Mx_bc3.ybc1 * uLe_i') + vec(Mx_bc3.ybc2 * uRi_i')
- yMx3 = Mx_bc3.Bbc * ybc3
- yMx = α * yMx - yMx3
- end
-
- # My
- # ybc = My_bc.ybc1 ⊗ vLo_i + My_bc.ybc2 ⊗ vUp_i
- ybc = vec(vLo_i * My_bc.ybc1') + vec(vUp_i * My_bc.ybc2')
- yMy = My_bc.Bbc * ybc
- if order4
- ybc3 = My_bc3.ybc1 ⊗ vLo_i + My_bc3.ybc2 ⊗ vUp_i
- yMy3 = My_bc3.Bbc * ybc3
- yMy = α * yMy - yMy3
- end
-
- yM = yMx + yMy
-
- # Time derivative of divergence
- if bc_unsteady
- ybc = dudtLe_i ⊗ Mx_bc.ybc1 + dudtRi_i ⊗ Mx_bc.ybc2
- ydMx = Mx_bc.Bbc * ybc
- if order4
- ybc3 = dudtLe_i ⊗ Mx_bc3.ybc1 + dudtRi_i ⊗ Mx_bc3.ybc2
- ydMx3 = Mx_bc3.Bbc * ybc3
- ydMx = α * ydMx - ydMx3
- end
-
- # My
- ybc = My_bc.ybc1 ⊗ dvdtLo_i + My_bc.ybc2 ⊗ dvdtUp_i
- ydMy = My_bc.Bbc * ybc
- if order4
- ybc3 = My_bc3.ybc1 ⊗ dvdtLo_i + My_bc3.ybc2 ⊗ dvdtUp_i
- ydMy3 = My_bc3.Bbc * ybc3
- ydMy = α * ydMy - ydMy3
- end
-
- ydM = ydMx + ydMy
- else
- ydM = zeros(T, Np)
- end
-
- ## Boundary conditions for pressure
-
- # Left and right side
- y1D_le = zeros(T, Nux_in)
- y1D_ri = zeros(T, Nux_in)
- boundary_conditions.u.x[1] == :pressure && (y1D_le[1] = -1)
- boundary_conditions.u.x[2] == :pressure && (y1D_ri[end] = 1)
- y_px = (hy .* p_bc.x[1]) ⊗ y1D_le + (hy .* p_bc.x[2]) ⊗ y1D_ri
-
- # Lower and upper side
- y1D_lo = zeros(T, Nvy_in)
- y1D_up = zeros(T, Nvy_in)
- boundary_conditions.v.y[1] == :pressure && (y1D_lo[1] = -1)
- boundary_conditions.v.y[2] == :pressure && (y1D_up[end] = 1)
- y_py = y1D_lo ⊗ (hx .* p_bc.y[1]) + y1D_up ⊗ (hx .* p_bc.y[2])
-
- y_p = [y_px; y_py]
-
- ## Boundary conditions for averaging
- # Au_ux
- ybc = uLe_i ⊗ Au_ux_bc.ybc1 + uRi_i ⊗ Au_ux_bc.ybc2
- yAu_ux = Au_ux_bc.Bbc * ybc
-
- # Au_uy
- ybc = Au_uy_bc.ybc1 ⊗ uLo_i + Au_uy_bc.ybc2 ⊗ uUp_i
- yAu_uy = Au_uy_bc.Bbc * ybc
-
- # Av_vx
- ybc = vLe_i ⊗ Av_vx_bc.ybc1 + vRi_i ⊗ Av_vx_bc.ybc2
- yAv_vx = Av_vx_bc.Bbc * ybc
-
- # Av_vy
- ybc = Av_vy_bc.ybc1 ⊗ vLo_i + Av_vy_bc.ybc2 ⊗ vUp_i
- yAv_vy = Av_vy_bc.Bbc * ybc
-
- if order4
- # Au_ux
- ybc3 = uLe_i ⊗ Au_ux_bc3.ybc1 + uRi_i ⊗ Au_ux_bc3.ybc2
- yAu_ux3 = Au_ux_bc3.Bbc * ybc3
-
- # Au_uy
- ybc3 = Au_uy_bc3.ybc1 ⊗ uLo_i + Au_uy_bc3.ybc2 ⊗ uUp_i
- yAu_uy3 = Au_uy_bc3.Bbc * ybc3
-
- # Av_vx
- ybc3 = vLe_i ⊗ Av_vx_bc3.ybc1 + vRi_i ⊗ Av_vx_bc3.ybc2
- yAv_vx3 = Av_vx_bc3.Bbc * ybc3
-
- # Av_vy
- ybc3 = Av_vy_bc3.ybc1 ⊗ vLo_i + Av_vy_bc3.ybc2 ⊗ vUp_i
- yAv_vy3 = Av_vy_bc3.Bbc * ybc3
- end
-
- ## Boundary conditions for diffusion
- if order4
- ybc1 = uLe_i ⊗ Su_ux_bc.ybc1 + uRi_i ⊗ Su_ux_bc.ybc2
- ybc3 = uLe_i ⊗ Su_ux_bc3.ybc1 + uRi_i ⊗ Su_ux_bc3.ybc2
- ySu_ux = α * Su_ux_bc.Bbc * ybc1 - Su_ux_bc3.Bbc * ybc3
-
- ybc1 = Su_uy_bc.ybc1 ⊗ uLo_i + Su_uy_bc.ybc2 ⊗ uUp_i
- ybc3 = Su_uy_bc3.ybc1 ⊗ uLo_i + Su_uy_bc3.ybc2 ⊗ uUp_i
- ySu_uy = α * Su_uy_bc.Bbc * ybc1 - Su_uy_bc3.Bbc * ybc3
-
- ybc1 = vLe_i ⊗ Sv_vx_bc.ybc1 + vRi_i ⊗ Sv_vx_bc.ybc2
- ybc3 = vLe_i ⊗ Sv_vx_bc3.ybc1 + vRi_i ⊗ Sv_vx_bc3.ybc2
- ySv_vx = α * Sv_vx_bc.Bbc * ybc1 - Sv_vx_bc3.Bbc * ybc3
-
- ybc1 = Sv_vy_bc.ybc1 ⊗ vLo_i + Sv_vy_bc.ybc2 ⊗ vUp_i
- ybc3 = Sv_vy_bc3.ybc1 ⊗ vLo_i + Sv_vy_bc3.ybc2 ⊗ vUp_i
- ySv_vy = α * Sv_vy_bc.Bbc * ybc1 - Sv_vy_bc3.Bbc * ybc3
-
- if viscosity_model isa LaminarModel
- yDiffu = Diffux_div * ySu_ux + Diffuy_div * ySu_uy
- yDiffv = Diffvx_div * ySv_vx + Diffvy_div * ySv_vy
- yDiff = [yDiffu; yDiffv]
- else
- error("fourth order turbulent diffusion not implemented")
- end
- else
- # Su_ux
- ybc = uLe_i ⊗ Su_ux_bc.ybc1 + uRi_i ⊗ Su_ux_bc.ybc2
- ySu_ux = Su_ux_bc.Bbc * ybc
-
- # Su_uy
- ybc = Su_uy_bc.ybc1 ⊗ uLo_i + Su_uy_bc.ybc2 ⊗ uUp_i
- ySu_uy = Su_uy_bc.Bbc * ybc
-
- Sv_uy = Sv_uy_bc_lr.B2D * Sv_uy_bc_lu.B2D
-
- # Sv_uy (left/right)
- ybc = vLe ⊗ Sv_uy_bc_lr.ybc1 + vRi ⊗ Sv_uy_bc_lr.ybc2
- ySv_uy_lr = Sv_uy_bc_lr.Bbc * ybc
-
- # Iv_uy (low/up)
- ybc = Sv_uy_bc_lu.ybc1 ⊗ vLo_i + Sv_uy_bc_lu.ybc2 ⊗ vUp_i
- ySv_uy_lu = Sv_uy_bc_lr.B2D * Sv_uy_bc_lu.Bbc * ybc
-
- ySv_uy = ySv_uy_lr + ySv_uy_lu
-
- # Su_vx (low/up)
- ybc = Su_vx_bc_lu.ybc1 ⊗ uLo + Su_vx_bc_lu.ybc2 ⊗ uUp
- ySu_vx_lu = Su_vx_bc_lu.Bbc * ybc
-
- # Su_vx (left/right)
- ybc = uLe_i ⊗ Su_vx_bc_lr.ybc1 + uRi_i ⊗ Su_vx_bc_lr.ybc2
- ySu_vx_lr = Su_vx_bc_lu.B2D * Su_vx_bc_lr.Bbc * ybc
- ySu_vx = ySu_vx_lr + ySu_vx_lu
-
- # Sv_vx
- ybc = vLe_i ⊗ Sv_vx_bc.ybc1 + vRi_i ⊗ Sv_vx_bc.ybc2
- ySv_vx = Sv_vx_bc.Bbc * ybc
-
- # Sv_vy
- ybc = Sv_vy_bc.ybc1 ⊗ vLo_i + Sv_vy_bc.ybc2 ⊗ vUp_i
- ySv_vy = Sv_vy_bc.Bbc * ybc
-
- yDiffu = Dux * ySu_ux + Duy * ySu_uy
- yDiffv = Dvx * ySv_vx + Dvy * ySv_vy
- yDiff = [yDiffu; yDiffv]
-
- Ωu⁻¹ = 1 ./ Ω[indu]
- Ωv⁻¹ = 1 ./ Ω[indv]
-
- # Diffusive terms in finite-difference setting, without viscosity
- yDiffu_f = Ωu⁻¹ .* (Dux * ySu_ux + Duy * ySu_uy)
- yDiffv_f = Ωv⁻¹ .* (Dvx * ySv_vx + Dvy * ySv_vy)
- end
-
- ## Boundary conditions for interpolation
-
- # Iu_ux
- ybc = uLe_i ⊗ Iu_ux_bc.ybc1 + uRi_i ⊗ Iu_ux_bc.ybc2
- yIu_ux = Iu_ux_bc.Bbc * ybc
- if order4
- ybc3 = uLe_i ⊗ Iu_ux_bc3.ybc1 + uRi_i ⊗ Iu_ux_bc3.ybc2
- yIu_ux3 = Iu_ux_bc3.Bbc * ybc3
- end
-
- # Iv_uy (left/right)
- ybc = vLe ⊗ Iv_uy_bc_lr.ybc1 + vRi ⊗ Iv_uy_bc_lr.ybc2
- yIv_uy_lr = Iv_uy_bc_lr.Bbc * ybc
-
- # Iv_uy (low/up)
- ybc = Iv_uy_bc_lu.ybc1 ⊗ vLo_i + Iv_uy_bc_lu.ybc2 ⊗ vUp_i
- yIv_uy_lu = Iv_uy_bc_lr.B2D * Iv_uy_bc_lu.Bbc * ybc
- yIv_uy = yIv_uy_lr + yIv_uy_lu
-
- if order4
- if boundary_conditions.v.y[1] == :dirichlet
- vLe_ext = [2 * vLe[1] - vLe[2]; vLe]
- vRi_ext = [2 * vRi[1] - vRi[2]; vRi]
- elseif boundary_conditions.v.y[1] == :periodic
- vLe_ext = [0; vLe]
- vRi_ext = [0; vRi]
- elseif boundary_conditions.v.y[1] == :pressure
- vLe_ext = [vLe[2]; vLe]
- vRi_ext = [vRi[2]; vRi]
- end
- if boundary_conditions.v.y[2] == :dirichlet
- vLe_ext = [vLe_ext; 2 * vLe[end] - vLe[end-1]]
- vRi_ext = [vRi_ext; 2 * vRi[1] - vRi[2]]
- elseif boundary_conditions.v.y[2] == :periodic
- vLe_ext = [vLe_ext; 0]
- vRi_ext = [vRi_ext; 0]
- elseif boundary_conditions.v.y[2] == :pressure
- vLe_ext = [vLe_ext; vLe[end-1]]
- vRi_ext = [vRi_ext; vRi[end-1]]
- end
- ybc3 = vLe_ext ⊗ Iv_uy_bc_lr3.ybc1 + vRi_ext ⊗ Iv_uy_bc_lr3.ybc2
- yIv_uy_lr3 = Iv_uy_bc_lr3.Bbc * ybc3
-
- ybc3 = Iv_uy_bc_lu3.ybc1 ⊗ vLo_i + Iv_uy_bc_lu3.ybc2 ⊗ vUp_i
- yIv_uy_lu3 = Iv_uy_bc_lr3.B2D * Iv_uy_bc_lu3.Bbc * ybc3
- yIv_uy3 = yIv_uy_lr3 + yIv_uy_lu3
- end
-
- # Iu_vx (low/up)
- ybc = Iu_vx_bc_lu.ybc1 ⊗ uLo + Iu_vx_bc_lu.ybc2 ⊗ uUp
- yIu_vx_lu = Iu_vx_bc_lu.Bbc * ybc
-
- # Iu_vx (left/right)
- ybc = uLe_i ⊗ Iu_vx_bc_lr.ybc1 + uRi_i ⊗ Iu_vx_bc_lr.ybc2
- yIu_vx_lr = Iu_vx_bc_lu.B2D * Iu_vx_bc_lr.Bbc * ybc
- yIu_vx = yIu_vx_lr + yIu_vx_lu
-
- if order4
- if boundary_conditions.u.x[1] == :dirichlet
- uLo_ext = [2 * uLo[1] - uLo[2]; uLo]
- uUp_ext = [2 * uUp[1] - uUp[2]; uUp]
- elseif boundary_conditions.u.x[1] == :periodic
- uLo_ext = [0; uLo]
- uUp_ext = [0; uUp]
- elseif boundary_conditions.u.x[1] == :pressure
- uLo_ext = [uLo[2]; uLo]
- uUp_ext = [uUp[2]; uUp]
- end
- if boundary_conditions.u.x[2] == :dirichlet
- uLo_ext = [uLo_ext; 2 * uLo[end] - uLo[end-1]]
- uUp_ext = [uUp_ext; 2 * uUp[1] - uUp[2]]
- elseif boundary_conditions.u.x[2] == :periodic
- uLo_ext = [uLo_ext; 0]
- uUp_ext = [uUp_ext; 0]
- elseif boundary_conditions.u.x[2] == :pressure
- uLo_ext = [uLo_ext; uLo[end-1]]
- uUp_ext = [uUp_ext; uUp[end-1]]
- end
- ybc3 = Iu_vx_bc_lu3.ybc1 ⊗ uLo_ext + Iu_vx_bc_lu3.ybc2 ⊗ uUp_ext
- yIu_vx_lu3 = Iu_vx_bc_lu3.Bbc * ybc3
-
- ybc3 = uLe_i ⊗ Iu_vx_bc_lr3.ybc1 + uRi_i ⊗ Iu_vx_bc_lr3.ybc2
- yIu_vx_lr3 = Iu_vx_bc_lu3.B2D * Iu_vx_bc_lr3.Bbc * ybc3
- yIu_vx3 = yIu_vx_lr3 + yIu_vx_lu3
- end
-
- # Iv_vy
- ybc = Iv_vy_bc.ybc1 ⊗ vLo_i + Iv_vy_bc.ybc2 ⊗ vUp_i
- yIv_vy = Iv_vy_bc.Bbc * ybc
- if order4
- ybc3 = Iv_vy_bc3.ybc1 ⊗ vLo_i + Iv_vy_bc3.ybc2 ⊗ vUp_i
- yIv_vy3 = Iv_vy_bc3.Bbc * ybc3
- end
-
- ## Group BC vectors
- bc_vectors = (;
- yM,
- ydM,
- y_p,
- yAu_ux,
- yAu_uy,
- yAv_vx,
- yAv_vy,
- yDiff,
- yIu_ux,
- yIv_uy,
- yIu_vx,
- yIv_vy,
- )
-
- if order4
- bc_vectors = (;
- bc_vectors...,
- yAu_ux3,
- yAu_uy3,
- yAv_vx3,
- yAv_vy3,
- yIu_ux3,
- yIv_uy3,
- yIu_vx3,
- yIv_vy3,
- )
- else
- # Use values directly (see diffusion.jl and strain_tensor.jl)
- bc_vectors = (;
- bc_vectors...,
- ySu_ux,
- ySu_uy,
- ySu_vx,
- ySv_vx,
- ySv_vy,
- ySv_uy,
- yDiffu_f,
- yDiffv_f,
- )
- end
-
- if viscosity_model isa Union{QRModel,SmagorinskyModel,MixingLengthModel}
- (; Aν_vy_bc) = operators
- (; Cux_k_bc, Cuy_k_bc, Cvx_k_bc, Cvy_k_bc, Auy_k_bc, Avx_k_bc) = operators
-
- # Set BC for turbulent viscosity nu_t
- # In the periodic case, the value of nu_t is not needed
- # In all other cases, homogeneous (zero) Neumann conditions are used
- nuLe = zeros(T, Npy)
- nuRi = zeros(T, Npy)
- nuLo = zeros(T, Npx)
- nuUp = zeros(T, Npx)
-
- ## Nu_ux
- (; Aν_ux_bc) = operators
- ybc = nuLe ⊗ Aν_ux_bc.ybc1 + nuRi ⊗ Aν_ux_bc.ybc2
- yAν_ux = Aν_ux_bc.Bbc * ybc
-
- ## Nu_uy
- (; Aν_uy_bc_lr, Aν_uy_bc_lu) = operators
-
- nuLe_i = [nuLe[1]; nuLe; nuLe[end]]
- nuRi_i = [nuRi[1]; nuRi; nuRi[end]]
-
- # In x-direction
- ybc = nuLe_i ⊗ Aν_uy_bc_lr.ybc1 + nuRi_i ⊗ Aν_uy_bc_lr.ybc2
- yAν_uy_lr = Aν_uy_bc_lr.B2D * ybc
-
- # In y-direction
- ybc = Aν_uy_bc_lu.ybc1 ⊗ nuLo + Aν_uy_bc_lu.ybc2 ⊗ nuUp
- yAν_uy_lu = Aν_uy_bc_lu.B2D * ybc
-
- yAν_uy = yAν_uy_lu + yAν_uy_lr
-
- ## Nu_vx
- (; Aν_vx_bc_lr, Aν_vx_bc_lu) = operators
-
- nuLo_i = [nuLo[1]; nuLo; nuLo[end]]
- nuUp_i = [nuUp[1]; nuUp; nuUp[end]]
-
- # In y-direction
- ybc = Aν_vx_bc_lu.ybc1 ⊗ nuLo_i + Aν_vx_bc_lu.ybc2 ⊗ nuUp_i
- yAν_vx_lu = Aν_vx_bc_lu.B2D * ybc
-
- # In x-direction
- ybc = nuLe ⊗ Aν_vx_bc_lr.ybc1 + nuRi ⊗ Aν_vx_bc_lr.ybc2
- yAν_vx_lr = Aν_vx_bc_lr.B2D * ybc
-
- yAν_vx = yAν_vx_lu + yAν_vx_lr
-
- ## Nu_vy
- ybc = Aν_vy_bc.ybc1 ⊗ nuLo + Aν_vy_bc.ybc2 ⊗ nuUp
- yAν_vy = Aν_vy_bc.Bbc * ybc
-
- # Set BC for getting du/dx, du/dy, dv/dx, dv/dy at cell centers
-
- uLo_p = u_bc.(xp, y[1], t)
- uUp_p = u_bc.(xp, y[end], t)
-
- vLe_p = v_bc.(x[1], yp, t)
- vRi_p = v_bc.(x[end], yp, t)
-
- ybc = uLe_i ⊗ Cux_k_bc.ybc1 + uRi_i ⊗ Cux_k_bc.ybc2
- yCux_k = Cux_k_bc.Bbc * ybc
-
- ybc = uLe_i ⊗ Auy_k_bc.ybc1 + uRi_i ⊗ Auy_k_bc.ybc2
- yAuy_k = Auy_k_bc.Bbc * ybc
-
- ybc = Cuy_k_bc.ybc1 ⊗ uLo_p + Cuy_k_bc.ybc2 ⊗ uUp_p
- yCuy_k = Cuy_k_bc.Bbc * ybc
-
- ybc = Avx_k_bc.ybc1 ⊗ vLo_i + Avx_k_bc.ybc2 ⊗ vUp_i
- yAvx_k = Avx_k_bc.Bbc * ybc
-
- ybc = vLe_p ⊗ Cvx_k_bc.ybc1 + vRi_p ⊗ Cvx_k_bc.ybc2
- yCvx_k = Cvx_k_bc.Bbc * ybc
-
- ybc = Cvy_k_bc.ybc1 ⊗ vLo_i + Cvy_k_bc.ybc2 ⊗ vUp_i
- yCvy_k = Cvy_k_bc.Bbc * ybc
-
- bc_vectors = (;
- bc_vectors...,
- yAν_ux,
- yAν_uy,
- yAν_vx,
- yAν_vy,
- yCux_k,
- yCuy_k,
- yCvx_k,
- yCvy_k,
- yAuy_k,
- yAvx_k,
- )
- end
-
- bc_vectors
-end
-
-# 3D version
-function get_bc_vectors(::Dimension{3}, setup, t)
- (; grid, operators, boundary_conditions, viscosity_model) = setup
-
- (; Re) = viscosity_model
-
- (; u_bc, v_bc, w_bc, dudt_bc, dvdt_bc, dwdt_bc) = boundary_conditions
- (; p_bc, bc_unsteady) = boundary_conditions
-
- (; Ω, indu, indv, indw) = grid
- (; Nz, Np, Npx, Npy, Npz) = grid
- (; Nux_in, Nux_b, Nux_t, Nuy_in, Nuy_b, Nuy_t, Nuz_in, Nuz_b, Nuz_t) = grid
- (; Nvx_in, Nvx_b, Nvx_t, Nvy_in, Nvy_b, Nvy_t, Nvz_in, Nuz_b, Nvz_t) = grid
- (; Nwx_in, Nwx_b, Nwx_t, Nwy_in, Nwy_b, Nwy_t, Nwz_in, Nwz_b, Nwz_t) = grid
- (; xin, yin, zin, x, y, z, hx, hy, hz, xp, yp, zp) = grid
-
- (; Dux, Duy, Duz, Dvx, Dvy, Dvz, Dwx, Dwy, Dwz) = operators
-
- (; Au_ux_bc, Au_uy_bc, Au_uz_bc) = operators
- (; Av_vx_bc, Av_vy_bc, Av_vz_bc) = operators
- (; Aw_wx_bc, Aw_wy_bc, Aw_wz_bc) = operators
-
- (; Su_ux_bc, Su_uy_bc, Su_uz_bc) = operators
- (; Sv_vx_bc, Sv_vy_bc, Sv_vz_bc) = operators
- (; Sw_wx_bc, Sw_wy_bc, Sw_wz_bc) = operators
-
- (; Mx_bc, My_bc, Mz_bc) = operators
-
- (; Iu_ux_bc, Iv_vy_bc, Iw_wz_bc) = operators
- (; Iv_uy_bc_lr, Iv_uy_bc_lu, Iw_uz_bc_lr, Iw_uz_bc_bf) = operators
- (; Iu_vx_bc_lr, Iu_vx_bc_lu, Iw_vz_bc_lu, Iw_vz_bc_bf) = operators
- (; Iu_wx_bc_lr, Iu_wx_bc_bf, Iv_wy_bc_lu, Iv_wy_bc_bf) = operators
-
- (; Sv_uy_bc_lr, Sv_uy_bc_lu, Sw_uz_bc_lr, Sw_uz_bc_bf) = operators
- (; Su_vx_bc_lr, Su_vx_bc_lu, Sw_vz_bc_lu, Sw_vz_bc_bf) = operators
- (; Su_wx_bc_lr, Su_wx_bc_bf, Sv_wy_bc_lu, Sv_wy_bc_bf) = operators
-
- T = typeof(t)
-
- # TODO: Split up function into allocating part (constructor?) and mutating `update!`
-
- ## Get bc values
- uLe_i = reshape(u_bc.(x[1], yp, zp', t), :)
- uRi_i = reshape(u_bc.(x[end], yp, zp', t), :)
- uLo_i = reshape(u_bc.(xin, y[1], zp', t), :)
- uUp_i = reshape(u_bc.(xin, y[end], zp', t), :)
- uLo_i2 = reshape(u_bc.(x, y[1], zp', t), :)
- uUp_i2 = reshape(u_bc.(x, y[end], zp', t), :)
- uBa_i = reshape(u_bc.(xin, yp', z[1], t), :)
- uFr_i = reshape(u_bc.(xin, yp', z[end], t), :)
- uBa_i2 = reshape(u_bc.(x, yp', z[1], t), :)
- uFr_i2 = reshape(u_bc.(x, yp', z[end], t), :)
-
- vLe_i = reshape(v_bc.(x[1], yin, zp', t), :)
- vRi_i = reshape(v_bc.(x[end], yin, zp', t), :)
- vLe_i2 = reshape(v_bc.(x[1], y, zp', t), :)
- vRi_i2 = reshape(v_bc.(x[end], y, zp', t), :)
- vLo_i = reshape(v_bc.(xp, y[1], zp', t), :)
- vUp_i = reshape(v_bc.(xp, y[end], zp', t), :)
- vBa_i = reshape(v_bc.(xp, yin', z[1], t), :)
- vFr_i = reshape(v_bc.(xp, yin', z[end], t), :)
- vBa_i2 = reshape(v_bc.(xp, y', z[1], t), :)
- vFr_i2 = reshape(v_bc.(xp, y', z[end], t), :)
-
- wLe_i = reshape(w_bc.(x[1], yp, zin', t), :)
- wRi_i = reshape(w_bc.(x[end], yp, zin', t), :)
- wLe_i2 = reshape(w_bc.(x[1], yp, z', t), :)
- wRi_i2 = reshape(w_bc.(x[end], yp, z', t), :)
- wLo_i = reshape(w_bc.(xp, y[1], zin', t), :)
- wUp_i = reshape(w_bc.(xp, y[end], zin', t), :)
- wLo_i2 = reshape(w_bc.(xp, y[1], z', t), :)
- wUp_i2 = reshape(w_bc.(xp, y[end], z', t), :)
- wBa_i = reshape(w_bc.(xp, yp', z[1], t), :)
- wFr_i = reshape(w_bc.(xp, yp', z[end], t), :)
-
- if bc_unsteady
- dudtLe_i = reshape(dudt_bc.(x[1], yp, zp', t), :)
- dudtRi_i = reshape(dudt_bc.(x[end], yp, zp', t), :)
- dvdtLo_i = reshape(dvdt_bc.(xp, y[1], zp', t), :)
- dvdtUp_i = reshape(dvdt_bc.(xp, y[end], zp', t), :)
- dwdtBa_i = reshape(dwdt_bc.(xp, yp', z[1], t), :)
- dwdtFr_i = reshape(dwdt_bc.(xp, yp', z[end], t), :)
- end
-
- ## Boundary conditions for divergence
-
- # Mx
- ybc = uLe_i ⊗ Mx_bc.ybc1 + uRi_i ⊗ Mx_bc.ybc2
- yMx = Mx_bc.Bbc * ybc
-
- # My
- ybc = vLo_i ⊗ My_bc.ybc1 + vUp_i ⊗ My_bc.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nvy_b, Nvx_in, Nvz_in), (2, 1, 3)), :)
- yMy = My_bc.Bbc * ybc
-
- # Mz
- ybc = Mz_bc.ybc1 ⊗ wBa_i + Mz_bc.ybc2 ⊗ wFr_i
- yMz = Mz_bc.Bbc * ybc
-
- yM = yMx + yMy + yMz
-
- # Time derivative of divergence
- if bc_unsteady
- ybc = dudtLe_i ⊗ Mx_bc.ybc1 + dudtRi_i ⊗ Mx_bc.ybc2
- ydMx = Mx_bc.Bbc * ybc
-
- # My - order of kron is not correct, so reshape
- ybc = dvdtLo_i ⊗ My_bc.ybc1 + dvdtUp_i ⊗ My_bc.ybc2
- ybc = reshape(ybc, Nvy_b, Nvx_in, Nvz_in)
- ybc = permutedims(ybc, (2, 1, 3))
- ybc = reshape(ybc, :)
- ydMy = My_bc.Bbc * ybc
-
- # Mz
- ybc = Mz_bc.ybc1 ⊗ dwdtBa_i + Mz_bc.ybc2 ⊗ dwdtFr_i
- ydMz = Mz_bc.Bbc * ybc
-
- ydM = ydMx + ydMy + ydMz
- else
- ydM = zeros(T, Np)
- end
-
- ## Boundary conditions for pressure
-
- # Left and right side
- y1D_le = zeros(T, Nux_in)
- y1D_ri = zeros(T, Nux_in)
- boundary_conditions.u.x[1] == :pressure && (y1D_le[1] = -1)
- boundary_conditions.u.x[2] == :pressure && (y1D_ri[end] = 1)
- y_px = (p_bc.x[1] .* (hz ⊗ hy)) ⊗ y1D_le + (p_bc.x[2] .* (hz ⊗ hy)) ⊗ y1D_ri
-
- # Lower and upper side (order of kron not correct, so reshape)
- y1D_lo = zeros(T, Nvy_in)
- y1D_up = zeros(T, Nvy_in)
- boundary_conditions.v.y[1] == :pressure && (y1D_lo[1] = -1)
- boundary_conditions.v.y[2] == :pressure && (y1D_up[end] = 1)
- y_py = (p_bc.y[1] .* (hz ⊗ hx)) ⊗ y1D_lo + (p_bc.y[2] .* (hz ⊗ hx)) ⊗ y1D_up
- y_py = reshape(permutedims(reshape(y_py, Nvy_in, Nvx_in, Nvz_in), (2, 1, 3)), :)
-
- # Back and front side
- y1D_ba = zeros(T, Nwz_in)
- y1D_fr = zeros(T, Nwz_in)
- boundary_conditions.w.z[1] == :pressure && (y1D_ba[1] = -1)
- boundary_conditions.w.z[2] == :pressure && (y1D_fr[end] = 1)
- y_pz = y1D_ba ⊗ (p_bc.z[1] .* (hy ⊗ hx)) + y1D_fr ⊗ (p_bc.z[2] .* (hy ⊗ hx))
-
- y_p = [y_px; y_py; y_pz]
-
- ## Boundary conditions for averaging
- # Au_ux
- ybc = uLe_i ⊗ Au_ux_bc.ybc1 + uRi_i ⊗ Au_ux_bc.ybc2
- yAu_ux = Au_ux_bc.Bbc * ybc
-
- # Au_uy (order of kron is not correct, so permute)
- ybc = Au_uy_bc.ybc1 ⊗ uLo_i + Au_uy_bc.ybc2 ⊗ uUp_i
- ybc = reshape(permutedims(reshape(ybc, Nuy_b, Nux_in, Nuz_in), (2, 3, 1)), :)
- yAu_uy = Au_uy_bc.Bbc * ybc
-
- # Au_uz
- ybc = Au_uz_bc.ybc1 ⊗ uBa_i + Au_uz_bc.ybc2 ⊗ uFr_i
- yAu_uz = Au_uz_bc.Bbc * ybc
-
- # Av_vx
- ybc = vLe_i ⊗ Av_vx_bc.ybc1 + vRi_i ⊗ Av_vx_bc.ybc2
- yAv_vx = Av_vx_bc.Bbc * ybc
-
- # Av_vy (order of kron is not correct, so permute)
- ybc = Av_vy_bc.ybc1 ⊗ vLo_i + Av_vy_bc.ybc2 ⊗ vUp_i
- ybc = reshape(permutedims(reshape(ybc, Nvy_b, Nvx_in, Nvz_in), (2, 3, 1)), :)
- yAv_vy = Av_vy_bc.Bbc * ybc
-
- # Av_vz
- ybc = Av_vz_bc.ybc1 ⊗ vBa_i + Av_vz_bc.ybc2 ⊗ vFr_i
- yAv_vz = Av_vz_bc.Bbc * ybc
-
- # Aw_wx
- ybc = wLe_i ⊗ Aw_wx_bc.ybc1 + wRi_i ⊗ Aw_wx_bc.ybc2
- yAw_wx = Aw_wx_bc.Bbc * ybc
-
- # Aw_wy (order of kron is not correct, so permute)
- ybc = wLo_i ⊗ Aw_wy_bc.ybc1 + wUp_i ⊗ Aw_wy_bc.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nwy_b, Nwx_in, Nwz_in), (2, 3, 1)), :)
- yAw_wy = Aw_wy_bc.Bbc * ybc
-
- # Aw_wz
- ybc = Aw_wz_bc.ybc1 ⊗ wBa_i + Aw_wz_bc.ybc2 ⊗ wFr_i
- yAw_wz = Aw_wz_bc.Bbc * ybc
-
- ## Bounary conditions for diffusion
-
- # Su_ux
- ybc = uLe_i ⊗ Su_ux_bc.ybc1 + uRi_i ⊗ Su_ux_bc.ybc2
- ySu_ux = Su_ux_bc.Bbc * ybc
-
- # Su_uy
- ybc = uLo_i ⊗ Su_uy_bc.ybc1 + uUp_i ⊗ Su_uy_bc.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nuy_b, Nux_in, Nuz_in), (2, 1, 3)), :)
- ySu_uy = Su_uy_bc.Bbc * ybc
-
- # Su_uz
- ybc = Su_uz_bc.ybc1 ⊗ uBa_i + Su_uz_bc.ybc2 ⊗ uFr_i
- ySu_uz = Su_uz_bc.Bbc * ybc
-
- # Sv_uy (left/right)
- ybc = vLe_i2 ⊗ Sv_uy_bc_lr.ybc1 + vRi_i2 ⊗ Sv_uy_bc_lr.ybc2
- ySv_uy_lr = Sv_uy_bc_lr.Bbc * ybc
-
- # Sv_uy (low/up) (order of kron is not correct, so permute)
- ybc = Sv_uy_bc_lu.ybc1 ⊗ vLo_i + Sv_uy_bc_lu.ybc2 ⊗ vUp_i
- Nb = Nuy_in + 1 - Nvy_in
- Nb != 0 && (ybc = reshape(permutedims(reshape(ybc, Nb, Nvx_in, Nvz_in), (2, 1, 3)), :))
- ySv_uy_lu = Sv_uy_bc_lr.B3D * (Sv_uy_bc_lu.Bbc * ybc)
-
- # Sv_uy
- ySv_uy = ySv_uy_lr + ySv_uy_lu
-
- # Sw_uz (left/right)
- ybc = wLe_i2 ⊗ Sw_uz_bc_lr.ybc1 + wRi_i2 ⊗ Sw_uz_bc_lr.ybc2
- ySw_uz_lr = Sw_uz_bc_lr.Bbc * ybc
-
- # Sw_uz (back/front)
- ybc = Sw_uz_bc_bf.ybc1 ⊗ wBa_i + Sw_uz_bc_bf.ybc2 ⊗ wFr_i
- ySw_uz_bf = Sw_uz_bc_lr.B3D * Sw_uz_bc_bf.Bbc * ybc
-
- # Sw_uz
- ySw_uz = ySw_uz_lr + ySw_uz_bf
-
- # Sv_vx
- ybc = vLe_i ⊗ Sv_vx_bc.ybc1 + vRi_i ⊗ Sv_vx_bc.ybc2
- ySv_vx = Sv_vx_bc.Bbc * ybc
-
- # Sv_vy (order of kron is not correct, so permute)
- ybc = vLo_i ⊗ Sv_vy_bc.ybc1 + vUp_i ⊗ Sv_vy_bc.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nvy_b, Nvx_in, Nvz_in), (2, 1, 3)), :)
- ySv_vy = Sv_vy_bc.Bbc * ybc
-
- # Sv_vz
- ybc = Sv_vz_bc.ybc1 ⊗ vBa_i + Sv_vz_bc.ybc2 ⊗ vFr_i
- ySv_vz = Sv_vz_bc.Bbc * ybc
-
- # Su_vx (low/up) (order of kron is not correct, so permute)
- ybc = uLo_i2 ⊗ Su_vx_bc_lu.ybc1 + uUp_i2 ⊗ Su_vx_bc_lu.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nuy_b, Nvx_t - 1, Nvz_in), (2, 1, 3)), :)
- ySu_vx_lu = Su_vx_bc_lu.Bbc * ybc
-
- # Su_vx (left/right)
- ybc = uLe_i ⊗ Su_vx_bc_lr.ybc1 + uRi_i ⊗ Su_vx_bc_lr.ybc2
- ySu_vx_lr = Su_vx_bc_lu.B3D * Su_vx_bc_lr.Bbc * ybc
-
- # Su_vx
- ySu_vx = ySu_vx_lr + ySu_vx_lu
-
- # Sw_vz (low/up) (order of kron is not correct, so permute)
- ybc = wLo_i2 ⊗ Sw_vz_bc_lu.ybc1 + wUp_i2 ⊗ Sw_vz_bc_lu.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nwy_b, Nwx_in, Nz + 1), (2, 1, 3)), :)
- ySw_vz_lu = Sw_vz_bc_lu.Bbc * ybc
-
- # Sw_vz (back/front)
- ybc = Sw_vz_bc_bf.ybc1 ⊗ wBa_i + Sw_vz_bc_bf.ybc2 ⊗ wFr_i
- ySw_vz_bf = Sw_vz_bc_lu.B3D * Sw_vz_bc_bf.Bbc * ybc
-
- # Sw_vz
- ySw_vz = ySw_vz_lu + ySw_vz_bf
-
- # Sw_wx
- ybc = wLe_i ⊗ Sw_wx_bc.ybc1 + wRi_i ⊗ Sw_wx_bc.ybc2
- ySw_wx = Sw_wx_bc.Bbc * ybc
-
- # Sw_wy (order of kron is not correct, so permute)
- ybc = wLo_i ⊗ Sw_wy_bc.ybc1 + wUp_i ⊗ Sw_wy_bc.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nwy_b, Nwx_in, Nwz_in), (2, 1, 3)), :)
- ySw_wy = Sw_wy_bc.Bbc * ybc
-
- # Sw_wz
- ybc = Sw_wz_bc.ybc1 ⊗ wBa_i + Sw_wz_bc.ybc2 ⊗ wFr_i
- ySw_wz = Sw_wz_bc.Bbc * ybc
-
- # Su_wx (back/front)
- ybc = Su_wx_bc_bf.ybc1 ⊗ uBa_i2 + Su_wx_bc_bf.ybc2 ⊗ uFr_i2
- ySu_wx_bf = Su_wx_bc_bf.Bbc * ybc
-
- # Su_wx (left/right)
- ybc = uLe_i ⊗ Su_wx_bc_lr.ybc1 + uRi_i ⊗ Su_wx_bc_lr.ybc2
- ySu_wx_lr = Su_wx_bc_bf.B3D * Su_wx_bc_lr.Bbc * ybc
-
- # Su_wx
- ySu_wx = ySu_wx_bf + ySu_wx_lr
-
- # Sv_wy (back/front)
- ybc = Sv_wy_bc_bf.ybc1 ⊗ vBa_i2 + Sv_wy_bc_bf.ybc2 ⊗ vFr_i2
- ySv_wy_bf = Sv_wy_bc_bf.Bbc * ybc
-
- # Sv_wy (low/up)
- ybc = vLo_i ⊗ Sv_wy_bc_lu.ybc1 + vUp_i ⊗ Sv_wy_bc_lu.ybc2
- Nb = Nwy_in + 1 - Nvy_in
- Nb != 0 && (ybc = reshape(permutedims(reshape(ybc, Nb, Nvx_in, Nvz_in), (2, 1, 3)), :))
- ySv_wy_lu = Sv_wy_bc_bf.B3D * Sv_wy_bc_lu.Bbc * ybc
-
- # Sv_wy
- ySv_wy = ySv_wy_bf + ySv_wy_lu
-
- yDiffu = Dux * ySu_ux + Duy * ySu_uy + Duz * ySu_uz
- yDiffv = Dvx * ySv_vx + Dvy * ySv_vy + Dvz * ySv_vz
- yDiffw = Dwx * ySw_wx + Dwy * ySw_wy + Dwz * ySw_wz
- yDiff = [yDiffu; yDiffv; yDiffw]
-
- Ωu⁻¹ = 1 ./ Ω[indu]
- Ωv⁻¹ = 1 ./ Ω[indv]
- Ωw⁻¹ = 1 ./ Ω[indw]
-
- # Diffusive terms in finite-difference setting, without viscosity
- yDiffu_f = Ωu⁻¹ .* (Dux * ySu_ux + Duy * ySu_uy + Duz * ySu_uz)
- yDiffv_f = Ωv⁻¹ .* (Dvx * ySv_vx + Dvy * ySv_vy + Dvz * ySv_vz)
- yDiffw_f = Ωw⁻¹ .* (Dwx * ySw_wx + Dwy * ySw_wy + Dwz * ySw_wz)
-
- ## Boundary conditions for interpolation
-
- # Iu_ux
- ybc = uLe_i ⊗ Iu_ux_bc.ybc1 + uRi_i ⊗ Iu_ux_bc.ybc2
- yIu_ux = Iu_ux_bc.Bbc * ybc
-
- # Iv_uy (left/right)
- ybc = vLe_i2 ⊗ Iv_uy_bc_lr.ybc1 + vRi_i2 ⊗ Iv_uy_bc_lr.ybc2
- yIv_uy_lr = Iv_uy_bc_lr.Bbc * ybc
-
- # Iv_uy (low/up) (order of kron is not correct, so permute)
- ybc = vLo_i ⊗ Iv_uy_bc_lu.ybc1 + vUp_i ⊗ Iv_uy_bc_lu.ybc2
- Nb = Nuy_in + 1 - Nvy_in
- Nb != 0 && (ybc = reshape(permutedims(reshape(ybc, Nb, Nvx_in, Nvz_in), (2, 1, 3)), :))
- yIv_uy_lu = Iv_uy_bc_lr.B3D * Iv_uy_bc_lu.Bbc * ybc
-
- yIv_uy = yIv_uy_lr + yIv_uy_lu
-
- # Iw_uz (left/right)
- ybc = wLe_i2 ⊗ Iw_uz_bc_lr.ybc1 + wRi_i2 ⊗ Iw_uz_bc_lr.ybc2
- yIw_uz_lr = Iw_uz_bc_lr.Bbc * ybc
-
- # Iw_uz (back/front)
- ybc = Iw_uz_bc_bf.ybc1 ⊗ wBa_i + Iw_uz_bc_bf.ybc2 ⊗ wFr_i
- yIw_uz_bf = Iw_uz_bc_lr.B3D * Iw_uz_bc_bf.Bbc * ybc
-
- # Iw_uz
- yIw_uz = yIw_uz_lr + yIw_uz_bf
-
- # Iu_vx (low/up) (order of kron is not correct, so permute)
- ybc = uLo_i2 ⊗ Iu_vx_bc_lu.ybc1 + uUp_i2 ⊗ Iu_vx_bc_lu.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nuy_b, Nvx_t - 1, Nvz_in), (2, 1, 3)), :)
- yIu_vx_lu = Iu_vx_bc_lu.Bbc * ybc
-
- # Iu_vx (left/right)
- ybc = uLe_i ⊗ Iu_vx_bc_lr.ybc1 + uRi_i ⊗ Iu_vx_bc_lr.ybc2
- yIu_vx_lr = Iu_vx_bc_lu.B3D * Iu_vx_bc_lr.Bbc * ybc
-
- # Iu_vx
- yIu_vx = yIu_vx_lr + yIu_vx_lu
-
- # Iv_vy (order of kron is not correct)
- ybc = vLo_i ⊗ Iv_vy_bc.ybc1 + vUp_i ⊗ Iv_vy_bc.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nvy_b, Nvx_in, Nvz_in), (2, 1, 3)), :)
- yIv_vy = Iv_vy_bc.Bbc * ybc
-
- # Iw_vz (low/up) (order of kron is not correct, so permute)
- ybc = wLo_i2 ⊗ Iw_vz_bc_lu.ybc1 + wUp_i2 ⊗ Iw_vz_bc_lu.ybc2
- ybc = reshape(permutedims(reshape(ybc, Nwy_b, Nwx_in, Nz + 1), (2, 1, 3)), :)
- yIw_vz_lu = Iw_vz_bc_lu.Bbc * ybc
-
- # Iw_vz (back/front)
- ybc = Iw_vz_bc_bf.ybc1 ⊗ wBa_i + Iw_vz_bc_bf.ybc2 ⊗ wFr_i
- yIw_vz_bf = Iw_vz_bc_lu.B3D * Iw_vz_bc_bf.Bbc * ybc
-
- # Iw_vz
- yIw_vz = yIw_vz_lu + yIw_vz_bf
-
- # Iu_wx (back/front)
- ybc = Iu_wx_bc_bf.ybc1 ⊗ uBa_i2 + Iu_wx_bc_bf.ybc2 ⊗ uFr_i2
- yIu_wx_bf = Iu_wx_bc_bf.Bbc * ybc
-
- # Iu_wx (left/right)
- ybc = uLe_i ⊗ Iu_wx_bc_lr.ybc1 + uRi_i ⊗ Iu_wx_bc_lr.ybc2
- yIu_wx_lr = Iu_wx_bc_bf.B3D * Iu_wx_bc_lr.Bbc * ybc
-
- yIu_wx = yIu_wx_bf + yIu_wx_lr
-
- # Iv_wy (back/front)
- ybc = Iv_wy_bc_bf.ybc1 ⊗ vBa_i2 + Iv_wy_bc_bf.ybc2 ⊗ vFr_i2
- yIv_wy_bf = Iv_wy_bc_bf.Bbc * ybc
-
- # Iv_wy (low/up)
- ybc = vLo_i ⊗ Iv_wy_bc_lu.ybc1 + vUp_i ⊗ Iv_wy_bc_lu.ybc2
- Nb = Nwy_in + 1 - Nvy_in
- Nb != 0 && (ybc = reshape(permutedims(reshape(ybc, Nb, Nvx_in, Nvz_in), (2, 1, 3)), :))
- yIv_wy_lu = Iv_wy_bc_bf.B3D * Iv_wy_bc_lu.Bbc * ybc
-
- # Iv_wy
- yIv_wy = yIv_wy_bf + yIv_wy_lu
-
- # Iw_wz
- ybc = Iw_wz_bc.ybc1 ⊗ wBa_i + Iw_wz_bc.ybc2 ⊗ wFr_i
- yIw_wz = Iw_wz_bc.Bbc * ybc
-
- # Group vectors
- bc_vectors = (;
- yM,
- ydM,
- y_p,
- yAu_ux,
- yAu_uy,
- yAu_uz,
- yAv_vx,
- yAv_vy,
- yAv_vz,
- yAw_wx,
- yAw_wy,
- yAw_wz,
- yDiff,
- ySu_ux,
- ySu_uy,
- ySu_uz,
- ySv_vx,
- ySv_vy,
- ySv_vz,
- ySw_wx,
- ySw_wy,
- ySw_wz,
- ySu_vx,
- ySu_wx,
- ySv_uy,
- ySv_wy,
- ySw_uz,
- ySw_vz,
- yDiffu_f,
- yDiffv_f,
- yDiffw_f,
- yIu_ux,
- yIv_uy,
- yIw_uz,
- yIu_vx,
- yIv_vy,
- yIw_vz,
- yIu_wx,
- yIv_wy,
- yIw_wz,
- )
-
- if viscosity_model isa Union{QRModel,SmagorinskyModel,MixingLengthModel}
- (; Aν_vy_bc) = operators
- (; Cux_k_bc, Cuy_k_bc, Cuz_k_bc) = operators
- (; Cvx_k_bc, Cvy_k_bc, Cvz_k_bc) = operators
- (; Cwx_k_bc, Cwy_k_bc, Cwz_k_bc) = operators
- (; Auy_k_bc, Avx_k_bc) = operators
- (; Auz_k_bc, Awx_k_bc) = operators
- (; Awy_k_bc, Avz_k_bc) = operators
-
- # Set bc for turbulent viscosity nu_t
- # In the periodic case, the value of nu_t is not needed
- # In all other cases, homogeneous (zero) Neumann conditions are used
- nuLe = zeros(T, Npy)
- nuRi = zeros(T, Npy)
- nuLo = zeros(T, Npx)
- nuUp = zeros(T, Npx)
-
- ## Nu_ux
- (; Aν_ux_bc) = operators
- ybc = nuLe ⊗ Aν_ux_bc.ybc1 + nuRi ⊗ Aν_ux_bc.ybc2
- yAν_ux = Aν_ux_bc.Bbc * ybc
-
- ## Nu_uy
- (; Aν_uy_bc_lr, Aν_uy_bc_lu) = operators
-
- nuLe_i = [nuLe[1]; nuLe; nuLe[end]]
- nuRi_i = [nuRi[1]; nuRi; nuRi[end]]
-
- # In x-direction
- ybc = nuLe_i ⊗ Aν_uy_bc_lr.ybc1 + nuRi_i ⊗ Aν_uy_bc_lr.ybc2
- yAν_uy_lr = Aν_uy_bc_lr.B3D * ybc
-
- # In y-direction
- ybc = Aν_uy_bc_lu.ybc1 ⊗ nuLo + Aν_uy_bc_lu.ybc2 ⊗ nuUp
- yAν_uy_lu = Aν_uy_bc_lu.B3D * ybc
-
- yAν_uy = yAν_uy_lu + yAν_uy_lr
-
- ## Nu_vx
- (; Aν_vx_bc_lr, Aν_vx_bc_lu) = operators
-
- nuLo_i = [nuLo[1]; nuLo; nuLo[end]]
- nuUp_i = [nuUp[1]; nuUp; nuUp[end]]
-
- # In y-direction
- ybc = Aν_vx_bc_lu.ybc1 ⊗ nuLo_i + Aν_vx_bc_lu.ybc2 ⊗ nuUp_i
- yAν_vx_lu = Aν_vx_bc_lu.B3D * ybc
-
- # In x-direction
- ybc = nuLe ⊗ Aν_vx_bc_lr.ybc1 + nuRi ⊗ Aν_vx_bc_lr.ybc2
- yAν_vx_lr = Aν_vx_bc_lr.B3D * ybc
-
- yAν_vx = yAν_vx_lu + yAν_vx_lr
-
- ## Nu_vy
- ybc = Aν_vy_bc.ybc1 ⊗ nuLo + Aν_vy_bc.ybc2 ⊗ nuUp
- yAν_vy = Aν_vy_bc.Bbc * ybc
-
- # Set bc for getting du/dx, du/dy, dv/dx, dv/dy at cell centers
-
- uLo_p = u_bc.(xp, y[1], t)
- uUp_p = u_bc.(xp, y[end], t)
-
- vLe_p = v_bc.(x[1], yp, t)
- vRi_p = v_bc.(x[end], yp, t)
-
- ybc = uLe_i ⊗ Cux_k_bc.ybc1 + uRi_i ⊗ Cux_k_bc.ybc2
- yCux_k = Cux_k_bc.Bbc * ybc
-
- ybc = uLe_i ⊗ Auy_k_bc.ybc1 + uRi_i ⊗ Auy_k_bc.ybc2
- yAuy_k = Auy_k_bc.Bbc * ybc
-
- ybc = Cuy_k_bc.ybc1 ⊗ uLo_p + Cuy_k_bc.ybc2 ⊗ uUp_p
- yCuy_k = Cuy_k_bc.Bbc * ybc
-
- ybc = Avx_k_bc.ybc1 ⊗ vLo_i + Avx_k_bc.ybc2 ⊗ vUp_i
- yAvx_k = Avx_k_bc.Bbc * ybc
-
- ybc = vLe_p ⊗ Cvx_k_bc.ybc1 + vRi_p ⊗ Cvx_k_bc.ybc2
- yCvx_k = Cvx_k_bc.Bbc * ybc
-
- ybc = Cvy_k_bc.ybc1 ⊗ vLo_i + Cvy_k_bc.ybc2 ⊗ vUp_i
- yCvy_k = Cvy_k_bc.Bbc * ybc
-
- bc_vectors = (;
- bc_vectors...,
- yAν_ux,
- yAν_uy,
- yAν_vx,
- yAν_vy,
- yCux_k,
- yCuy_k,
- yCvx_k,
- yCvy_k,
- yAuy_k,
- yAvx_k,
- )
- end
-
- bc_vectors
-end
diff --git a/src/create_initial_conditions.jl b/src/create_initial_conditions.jl
new file mode 100644
index 000000000..ed0a81129
--- /dev/null
+++ b/src/create_initial_conditions.jl
@@ -0,0 +1,215 @@
+"""
+ create_initial_conditions(
+ setup,
+ initial_velocity,
+ t = 0;
+ psolver = default_psolver(setup),
+ doproject = true,
+ )
+
+Create divergence free initial velocity field `u` at starting time `t`.
+The initial conditions of `u[α]` are specified by the function
+`initial_velocity(Dimension(α), x...)`.
+"""
+function create_initial_conditions(
+ setup,
+ initial_velocity,
+ t = convert(eltype(setup.grid.x[1]), 0);
+ psolver = default_psolver(setup),
+ doproject = true,
+)
+ (; grid) = setup
+ (; dimension, N, Iu, Ip, x, xp, Ω) = grid
+
+ T = eltype(x[1])
+ D = dimension()
+
+ # Allocate velocity
+ u = ntuple(d -> fill!(similar(x[1], N), 0), D)
+
+ # Initial velocities
+ for α = 1:D
+ xin = ntuple(
+ β -> reshape(α == β ? x[β][2:end] : xp[β], ntuple(Returns(1), β - 1)..., :),
+ D,
+ )
+ u[α][Iu[α]] .= initial_velocity.((Dimension(α),), xin...)[Iu[α]]
+ end
+
+ # Make velocity field divergence free
+ apply_bc_u!(u, t, setup)
+ if doproject
+ u = project(u, setup; psolver)
+ apply_bc_u!(u, t, setup)
+ end
+
+ # Initial conditions, including initial boundary conditions
+ u
+end
+
+# function create_spectrum(; setup, A, σ, s, rng = Random.default_rng())
+# (; dimension, x, N) = setup.grid
+# T = eltype(x[1])
+# D = dimension()
+# K = N .÷ 2
+# k = ntuple(
+# α -> reshape(1:K[α], ntuple(Returns(1), α - 1)..., :, ntuple(Returns(1), D - α)...),
+# D,
+# )
+# a = fill!(similar(x[1], Complex{T}, K), 1)
+# τ = T(2π)
+# a .*= prod(N) * A / sqrt(τ^2 * 2σ^2)
+# for α = 1:D
+# kα = k[α]
+# @. a *= exp(-max(abs(kα) - s, 0)^2 / 2σ^2)
+# end
+# @. a *= randn(rng, T) * exp(im * τ * rand(rng, T))
+# for α = 1:D
+# a = cat(a, reverse(a; dims = α); dims = α)
+# end
+# a
+# end
+
+function create_spectrum(; setup, kp, rng = Random.default_rng())
+ (; dimension, x, N) = setup.grid
+ T = eltype(x[1])
+ D = dimension()
+ τ = T(2π)
+
+ # Maximum wavenumber (remove ghost volumes)
+ K = @. (N - 2) ÷ 2
+
+ # Wavenumber vectors
+ kk = ntuple(
+ α -> reshape(
+ 0:K[α]-1,
+ ntuple(Returns(1), α - 1)...,
+ :,
+ ntuple(Returns(1), D - α)...,
+ ),
+ D,
+ )
+
+ # Wavevector magnitude
+ k = fill!(similar(x[1], K), 0)
+ for α = 1:D
+ @. k += kk[α]^2
+ end
+ k .= sqrt.(k)
+
+ # Shared magnitude
+ A = T(8τ / 3) / kp^5
+
+ # Velocity magnitude
+ # a = @. complex(1) * sqrt(A * k^4 * exp(-(k / kp)^2))
+ a = @. complex(1) * sqrt(A * k^4 * exp(-τ * (k / kp)^2))
+ a .*= prod(N)
+
+ # Apply random phase shift
+ ξ = ntuple(α -> rand!(rng, similar(x[1], K)), D)
+ for α = 1:D
+ a = cat(a, reverse(a; dims = α); dims = α)
+ ξ = ntuple(D) do β
+ s = α == β ? -1 : 1
+ ξβ = ξ[β]
+ cat(ξβ, reverse(s .* ξβ; dims = α); dims = α)
+ end
+ end
+ ξ = sum(ξ)
+ a = @. exp(im * τ * ξ) * a
+
+ KK = 2 .* K
+ kkkk = ntuple(
+ α -> reshape(
+ 0:KK[α]-1,
+ ntuple(Returns(1), α - 1)...,
+ :,
+ ntuple(Returns(1), D - α)...,
+ ),
+ D,
+ )
+ knorm = fill!(similar(x[1], KK), 0)
+ for α = 1:D
+ @. knorm += kkkk[α]^2
+ end
+ knorm .= sqrt.(knorm)
+
+ # Create random unit vector for each wavenumber
+ if D == 2
+ θ = rand!(rng, similar(x[1], KK))
+ e = (cospi.(2 .* θ), sinpi.(2 .* θ))
+ elseif D == 3
+ θ = rand!(rng, similar(x[1], KK))
+ ϕ = rand!(rng, similar(x[1], KK))
+ e = (sinpi.(θ) .* cospi.(2 .* ϕ), sinpi.(θ) .* sinpi.(2 .* ϕ), cospi.(θ))
+ end
+
+ # Remove non-divergence free part: (I - k k^T / k^2) e
+ ke = sum(α -> e[α] .* kkkk[α], 1:D)
+ for α = 1:D
+ e0 = e[α][1:1] # CUDA doesn't like e[α][1]
+ @. e[α] -= kkkk[α] * ke / knorm^2
+ # Restore k=0 component, which is divergence free anyways
+ e[α][1:1] .= e0
+ end
+
+ # Normalize
+ enorm = sqrt.(sum(α -> e[α] .^ 2, 1:D))
+ for α = 1:D
+ e[α] ./= enorm
+ end
+
+ # Split velocity magnitude a into velocity components a*eα
+ uhat = ntuple(D) do α
+ eα = e[α]
+ # for β = 1:D
+ # eα = cat(eα, reverse(eα; dims = β); dims = β)
+ # end
+ a .* eα
+ end
+end
+
+"""
+ random_field(
+ setup, t = 0;
+ A = 1,
+ kp = 10,
+ psolver = default_psolver(setup),
+ rng = Random.default_rng(),
+ )
+
+Create random field, as in [Orlandi2000](@cite).
+
+- `K`: Maximum wavenumber
+- `A`: Eddy amplitude scaling
+- `kp`: Peak energy wavenumber
+"""
+function random_field(
+ setup,
+ t = zero(eltype(setup.grid.x[1]));
+ A = 1,
+ kp = 10,
+ psolver = default_psolver(setup),
+ rng = Random.default_rng(),
+)
+ (; dimension, x, Ip, Ω) = setup.grid
+ D = dimension()
+ T = eltype(x[1])
+
+ # Create random velocity field
+ uhat = create_spectrum(; setup, kp, rng)
+ u = ifft.(uhat)
+ u = map(u -> A .* real.(u), u)
+
+ # Add ghost volumes (one on each side for periodic)
+ u = pad_circular.(u, 1; dims = 1:D)
+
+ # # Interpolate to staggered grid
+ # interpolate_p_u!(u, setup)
+
+ # Make velocity field divergence free on staggered grid
+ # (it is already diergence free on the "spectral grid")
+ apply_bc_u!(u, t, setup)
+ u = project(u, setup; psolver)
+ apply_bc_u!(u, t, setup)
+end
diff --git a/src/force/force.jl b/src/force/force.jl
deleted file mode 100644
index 5a4a7c111..000000000
--- a/src/force/force.jl
+++ /dev/null
@@ -1,28 +0,0 @@
-"""
- SteadyBodyForce(fu, fv, grid)
-
-Two-dimensional steady body force `f(x, y) = [fu(x, y), fv(x, y)]`.
-"""
-function SteadyBodyForce(fu, fv, grid)
- (; NV, indu, indv, xu, yu, xv, yv) = grid
- T = eltype(xu)
- F = zeros(T, NV)
- F[indu] .= reshape(fu.(xu, yu), :)
- F[indv] .= reshape(fv.(xv, yv), :)
- F
-end
-
-"""
- SteadyBodyForce(fu, fv, fw, grid)
-
-Three-dimensional steady body force `f(x, y, z) = [fu(x, y, z), fv(x, y, z), fw(x, y, z)]`.
-"""
-function SteadyBodyForce(fu, fv, fw, grid)
- (; NV, indu, indv, indw, xu, yu, zu, xv, yv, zv, xw, yw, zw) = grid
- T = eltype(xu)
- F = zeros(T, NV)
- F[indu] .= reshape(fu.(xu, yu, zu), :)
- F[indv] .= reshape(fv.(xv, yv, zv), :)
- F[indw] .= reshape(fw.(xw, yw, zw), :)
- F
-end
diff --git a/src/grid.jl b/src/grid.jl
new file mode 100644
index 000000000..f47e796ee
--- /dev/null
+++ b/src/grid.jl
@@ -0,0 +1,227 @@
+"""
+ Dimension(N)
+
+Represent an `N`-dimensional space.
+Returns `N` when called.
+
+```example
+julia> d = Dimension(3)
+Dimension{3}()
+
+julia> d()
+3
+```
+"""
+struct Dimension{N} end
+
+Dimension(N) = Dimension{N}()
+
+(::Dimension{N})() where {N} = N
+
+"""
+ max_size(grid)
+
+Get size of the largest grid element.
+"""
+function max_size(grid)
+ (; Δ) = grid
+ m = maximum.(Δ)
+ sqrt(sum(m .^ 2))
+end
+
+"""
+ cosine_grid(a, b, N)
+
+Create a nonuniform grid of `N + 1` points from `a` to `b` using a cosine
+profile, i.e.
+
+```math
+x_i = a + \\frac{1}{2} \\left( 1 - \\cos \\left( \\pi \\frac{i}{n} \\right) \\right)
+(b - a), \\quad i = 0, \\dots, N
+```
+
+See also [`stretched_grid`](@ref).
+"""
+function cosine_grid(a, b, N)
+ T = typeof(a)
+ i = T.(0:N)
+ @. a + (b - a) * (1 - cospi(i / N)) / 2
+end
+
+"""
+ stretched_grid(a, b, N, s = 1)
+
+Create a nonuniform grid of `N + 1` points from `a` to `b` with a stretch
+factor of `s`. If `s = 1`, return a uniform spacing from `a` to `b`. Otherwise,
+return a vector ``x \\in \\mathbb{R}^{N + 1}`` such that ``x_n = a + \\sum_{i =
+1}^n s^{i - 1} h`` for ``n = 0, \\dots , N``. Setting ``x_N = b`` then gives
+``h = (b - a) \\frac{1 - s}{1 - s^N}``, resulting in
+
+```math
+x_n = a + (b - a) \\frac{1 - s^n}{1 - s^N}, \\quad n = 0, \\dots, N.
+```
+
+Note that `stretched_grid(a, b, N, s)[n]` corresponds to ``x_{n - 1}``.
+
+See also [`cosine_grid`](@ref).
+"""
+function stretched_grid(a, b, N, s = 1)
+ s > 0 || error("The stretch factor must be positive")
+ if s ≈ 1
+ LinRange(a, b, N + 1)
+ else
+ map(i -> a + (b - a) * (1 - s^i) / (1 - s^N), 0:N)
+ end
+end
+
+"""
+ Grid(x, boundary_conditions)
+
+Create nonuniform Cartesian box mesh `x[1]` × ... × `x[d]` with boundary
+conditions `boundary_conditions`.
+"""
+function Grid(x, boundary_conditions; ArrayType = Array)
+ # Kill all LinRanges etc.
+ x = Array.(x)
+ xlims = extrema.(x)
+
+ D = length(x)
+ dimension = Dimension(D)
+
+ T = eltype(x[1])
+
+ # Add offset positions for ghost volumes
+ # For all BC, there is one ghost volume on each side,
+ # but not all of the ``d + 1`` fields have a component inside this ghost
+ # volume.
+ for d = 1:D
+ a, b = boundary_conditions[d]
+ ghost_a!(a, x[d])
+ ghost_b!(b, x[d])
+ end
+
+ # Number of finite volumes in each dimension, including ghost volumes
+ N = length.(x) .- 1
+
+ # Number of velocity DOFs in each dimension
+ Nu = ntuple(D) do α
+ ntuple(D) do β
+ na = offset_u(boundary_conditions[β][1], α == β, false)
+ nb = offset_u(boundary_conditions[β][2], α == β, true)
+ N[β] - na - nb
+ end
+ end
+
+ # Cartesian index ranges of velocity DOFs
+ Iu = ntuple(D) do α
+ Iuα = ntuple(D) do β
+ na = offset_u(boundary_conditions[β][1], α == β, false)
+ nb = offset_u(boundary_conditions[β][2], α == β, true)
+ 1+na:N[β]-nb
+ end
+ CartesianIndices(Iuα)
+ end
+
+ # Number of p DOFs in each dimension
+ Np = ntuple(D) do α
+ na = offset_p(boundary_conditions[α][1], false)
+ nb = offset_p(boundary_conditions[α][2], true)
+ N[α] - na - nb
+ end
+
+ # Cartesian index range of pressure DOFs
+ Ip = CartesianIndices(ntuple(D) do α
+ na = offset_p(boundary_conditions[α][1], false)
+ nb = offset_p(boundary_conditions[α][2], true)
+ 1+na:N[α]-nb
+ end)
+
+ xp = ntuple(d -> (x[d][1:end-1] .+ x[d][2:end]) ./ 2, D)
+
+ # Volume widths
+ # Infinitely thin widths are set to `eps(T)` to avoid division by zero
+ Δ = ntuple(D) do d
+ Δ = diff(x[d])
+ Δ[Δ.==0] .= eps(eltype(Δ))
+ Δ
+ end
+ Δu = ntuple(D) do d
+ Δu = push!(diff(xp[d]), Δ[d][end] / 2)
+ Δu[Δu.==0] .= eps(eltype(Δu))
+ Δu
+ end
+
+ # Reference volume sizes
+ Ω = ones(T, N...)
+ for d = 1:D
+ Ω .*= reshape(Δ[d], ntuple(Returns(1), d - 1)..., :)
+ end
+
+ # # Velocity volume sizes
+ # Ωu = ntuple(α -> ones(T, N), D)
+ # for α = 1:D, β = 1:D
+ # Ωu[α] .*= reshape((α == β ? Δu : Δ)[β], ntuple(Returns(1), β - 1)..., :)
+ # end
+
+ # # Vorticity volume sizes
+ # Ωω = ones(T, N)
+ # for α = 1:D
+ # Ωω .*= reshape(Δu[α], ntuple(Returns(1), α - 1)..., :)
+ # end
+
+ # # Velocity volume mid-sections
+ # Γu = ntuple(α -> ntuple(β -> ones(T, N), D), D)
+ # for α = 1:D, β = 1:D, γ in ((1:β-1)..., (β+1:D)...)
+ # Γu[α][β] .*=
+ # reshape(γ == β ? 1 : γ == α ? Δu[γ] : Δ[γ], ntuple(Returns(1), γ - 1)..., :)
+ # end
+
+ # # Velocity points
+ # Xu = ntuple(α -> ones(T, N))
+
+ # Interpolation weights from α-face centers x_I to x_{I + δ(β) / 2}
+ A = ntuple(
+ α -> ntuple(
+ β -> begin
+ if α == β
+ # Interpolation from face center to volume center
+ Aαβ1 = fill(T(1 / 2), N[α])
+ Aαβ1[1] = 1
+ Aαβ2 = fill(T(1 / 2), N[α])
+ Aαβ2[end] = 1
+ else
+ # Interpolation from α-face center to left (1) or right (2) α-face β-edge
+ # Aαβ1 = [(x[β][i] - xp[β][i-1]) / Δu[β][i-1] for i = 2:N[β]]
+ # Aαβ2 = 1 .- Aαβ1
+ Aαβ2 = [(x[β][i] - xp[β][i-1]) / Δu[β][i-1] for i = 2:N[β]]
+ Aαβ1 = 1 .- Aαβ2
+ pushfirst!(Aαβ1, 1)
+ push!(Aαβ2, 1)
+ end
+ (ArrayType(Aαβ1), ArrayType(Aαβ2))
+ end,
+ D,
+ ),
+ D,
+ )
+
+ # Grid quantities
+ (;
+ dimension,
+ N,
+ Nu,
+ Np,
+ Iu,
+ Ip,
+ xlims,
+ x = ArrayType.(x),
+ xp = ArrayType.(xp),
+ Δ = ArrayType.(Δ),
+ Δu = ArrayType.(Δu),
+ Ω = ArrayType(Ω),
+ # Ωu = ArrayType.(Ωu),
+ # Ωω = ArrayType(Ωω),
+ # Γu = ArrayType.(Γu),
+ A,
+ )
+end
diff --git a/src/grid/cosine_grid.jl b/src/grid/cosine_grid.jl
deleted file mode 100644
index 7f2fd4b97..000000000
--- a/src/grid/cosine_grid.jl
+++ /dev/null
@@ -1,18 +0,0 @@
-"""
- cosine_grid(a, b, N)
-
-Create a nonuniform grid of `N + 1` points from `a` to `b` using a cosine
-profile, i.e.
-
-```math
-x_i = a + \\frac{1}{2} \\left( 1 - \\cos \\left( \\pi \\frac{i}{n} \\right) \\right)
-(b - a), \\quad i = 0, \\dots, N
-```
-
-See also [`stretched_grid`](@ref).
-"""
-function cosine_grid(a, b, N)
- T = typeof(a)
- i = T.(0:N)
- @. a + (b - a) * (1 - cospi(i / N)) / 2
-end
diff --git a/src/grid/dimension.jl b/src/grid/dimension.jl
deleted file mode 100644
index 3cc8e8b81..000000000
--- a/src/grid/dimension.jl
+++ /dev/null
@@ -1,19 +0,0 @@
-"""
- Dimension(N)
-
-Represent an `N`-dimensional space.
-Returns `N` when called.
-
-```example
-julia> d = Dimension(3)
-Dimension{3}()
-
-julia> d()
-3
-```
-"""
-struct Dimension{N} end
-
-Dimension(N) = Dimension{N}()
-
-(::Dimension{N})() where {N} = N
diff --git a/src/grid/grid.jl b/src/grid/grid.jl
deleted file mode 100644
index aadcb43db..000000000
--- a/src/grid/grid.jl
+++ /dev/null
@@ -1,878 +0,0 @@
-"""
- Grid(x, y; boundary_conditions, order4 = false)
-
-Create nonuniform Cartesian box mesh `x` × `y` with boundary conditions `boundary_conditions`.
-If `order4` is `true`, a fourth order mesh is created.
-"""
-function Grid(x, y; boundary_conditions, order4 = false)
- dimension = Dimension(2)
-
- T = eltype(x)
-
- α = 81
- β = T(9 // 8)
-
- Nx = length(x) - 1
- Ny = length(y) - 1
- xlims = (x[1], x[end])
- ylims = (y[1], y[end])
-
- # Pressure positions
- xp = (x[1:(end-1)] + x[2:end]) / 2
- yp = (y[1:(end-1)] + y[2:end]) / 2
-
- # Distance between velocity points
- hx = diff(x)
- hy = diff(y)
-
- # Distance between pressure points
- gx = zeros(T, Nx + 1)
- gx[1] = hx[1] / 2
- gx[2:Nx] = (hx[1:(Nx-1)] + hx[2:Nx]) / 2
- gx[Nx+1] = hx[end] / 2
-
- gy = zeros(T, Ny + 1)
- gy[1] = hy[1] / 2
- gy[2:Ny] = (hy[1:(Ny-1)] + hy[2:Ny]) / 2
- gy[Ny+1] = hy[end] / 2
-
- # Number of pressure points
- Npx = Nx
- Npy = Ny
- Np = Npx * Npy
-
- ## u-volumes
- # x[1] x[2] x[3] .... x[Nx] x[Nx+1]
- # | | | | |
- # | | | | |
- # Dirichlet BC:
- # ULe u[1] u[2] .... u(Nx-1) uRi
- # Periodic BC:
- # u[1] u[2] u[3] .... u[Nx] u[1]
- # Pressure BC:
- # u[1] u[2] u[3] .... u[Nx] u[Nx+1]
-
- # x-dir
- Nux_b = 2 # Boundary points
- Nux_in = Nx + 1 # Inner points
- Nux_in -= boundary_conditions.u.x[1] ∈ [:dirichlet, :symmetric]
- Nux_in -= boundary_conditions.u.x[2] ∈ [:dirichlet, :symmetric]
- Nux_in -= boundary_conditions.u.x == (:periodic, :periodic)
- Nux_t = Nux_in + Nux_b # Total number
-
- # Y-dir
- Nuy_b = 2 # Boundary points
- Nuy_in = Ny # Inner points
- Nuy_t = Nuy_in + Nuy_b # Total number
-
- # Total number
- Nu = Nux_in * Nuy_in
-
- ## v-volumes
-
- # X-dir
- Nvx_b = 2 # Boundary points
- Nvx_in = Nx # Inner points
- Nvx_t = Nvx_in + Nvx_b # Total number
-
- # Y-dir
- Nvy_b = 2 # Boundary points
- Nvy_in = Ny + 1 # Inner points
- Nvy_in -= boundary_conditions.v.y[1] ∈ [:dirichlet, :symmetric]
- Nvy_in -= boundary_conditions.v.y[2] ∈ [:dirichlet, :symmetric]
- Nvy_in -= boundary_conditions.v.y == (:periodic, :periodic)
- Nvy_t = Nvy_in + Nvy_b # Total number
-
- # Total number
- Nv = Nvx_in * Nvy_in
-
- # Total number of velocity points
- NV = Nu + Nv
-
- ## For a grid with three times larger volumes:
- if order4
- hx3 = zeros(T, Nx)
- hx3[2:(end-1)] = hx[1:(end-2)] + hx[2:(end-1)] + hx[3:end]
- if boundary_conditions.u.x[1] == :periodic &&
- boundary_conditions.u.x[2] == :periodic
- hx3[1] = hx[end] + hx[1] + hx[2]
- hx3[end] = hx[end-1] + hx[end] + hx[1]
- else
- hx3[1] = 2 * hx[1] + hx[2]
- hx3[end] = hx[end-1] + 2 * hx[end]
- end
-
- hy3 = zeros(T, Ny)
- hy3[2:(end-1)] = hy[1:(end-2)] + hy[2:(end-1)] + hy[3:end]
- if boundary_conditions.v.y[1] == :periodic &&
- boundary_conditions.v.y[2] == :periodic
- hy3[1] = hy[end] + hy[1] + hy[2]
- hy3[end] = hy[end-1] + hy[end] + hy[1]
- else
- hy3[1] = 2 * hy[1] + hy[2]
- hy3[end] = hy[end-1] + 2 * hy[end]
- end
-
- hxi3 = copy(hx3)
- hyi3 = copy(hy3)
-
- # Distance between pressure points
- gx3 = zeros(T, Nx + 1)
- gx3[3:(Nx-1)] = gx[2:(end-3)] + gx[3:(end-2)] + gx[4:(end-1)]
- if boundary_conditions.u.x[1] == :periodic &&
- boundary_conditions.u.x[2] == :periodic
- gx3[1] = gx[end-1] + gx[end] + gx[1] + gx[2]
- gx3[2] = gx[end] + gx[1] + gx[2] + gx[3]
- gx3[end-1] = gx[end-2] + gx[end-1] + gx[end] + gx[1]
- gx3[end] = gx[end-1] + gx[end] + gx[1] + gx[2]
- else
- gx3[1] = 2 * gx[1] + 2 * gx[2]
- gx3[2] = 2 * gx[1] + gx[2] + gx[3]
- gx3[end-1] = 2 * gx[end] + gx[end-1] + gx[end-2]
- gx3[end] = 2 * gx[end] + 2 * gx[end-1]
- end
-
- # Distance between pressure points
- gy3 = zeros(T, Ny + 1)
- gy3[3:(Ny-1)] = gy[2:(end-3)] + gy[3:(end-2)] + gy[4:(end-1)]
- if boundary_conditions.v.y[1] == :periodic &&
- boundary_conditions.v.y[2] == :periodic
- gy3[1] = gy[end-1] + gy[end] + gy[1] + gy[2]
- gy3[2] = gy[end] + gy[1] + gy[2] + gy[3]
- gy3[end-1] = gy[end-2] + gy[end-1] + gy[end] + gy[1]
- gy3[end] = gy[end-1] + gy[end] + gy[1] + gy[2]
- else
- gy3[1] = 2 * gy[1] + 2 * gy[2]
- gy3[2] = 2 * gy[1] + gy[2] + gy[3]
- gy3[end-1] = 2 * gy[end] + gy[end-1] + gy[end-2]
- gy3[end] = 2 * gy[end] + 2 * gy[end-1]
- end
- end
-
- ## Adapt mesh metrics depending on number of volumes
-
- ## X-direction
-
- # gxd: differentiation
- gxd = copy(gx)
- gxd[1] = hx[1]
- gxd[end] = hx[end]
-
- # hxi: integration and hxd: differentiation
- # Map to find suitable size
- hxi = copy(hx)
-
- # Restrict Nx+2 to Nux_in+1 points
- if boundary_conditions.u.x == (:dirichlet, :dirichlet)
- xin = x[2:(end-1)]
- hxd = copy(hx)
- gxi = gx[2:(end-1)]
- diagpos = 1
-
- if order4
- hxd3 = [hx3[1]; hx3; hx3[end]]
- hxd13 = [hx[1]; hx; hx[end]]
- gxd3 = [2 * gx[1] + gx[2] + gx[3]; gx3; 2 * gx[end] + gx[end-1] + gx[end-2]]
- gxd13 = [gx[2]; 2 * gx[1]; gx[2:(end-1)]; 2 * gx[end]; gx[end-1]]
- gxi3 = gx3[2:(end-1)]
- end
- elseif boundary_conditions.u.x == (:dirichlet, :pressure)
- xin = x[2:end]
- hxd = [hx; hx[end]]
- gxi = gx[2:end]
- diagpos = 1
- elseif boundary_conditions.u.x == (:pressure, :dirichlet)
- xin = x[1:(end-1)]
- hxd = [hx[1]; hx]
- gxi = gx[1:(end-1)]
- diagpos = 0
- elseif boundary_conditions.u.x == (:pressure, :pressure)
- xin = x[1:end]
- hxd = [hx[1]; hx; hx[end]]
- gxi = copy(gx)
- diagpos = 0
- elseif boundary_conditions.u.x == (:periodic, :periodic)
- xin = x[1:(end-1)]
- hxd = [hx[end]; hx]
- gxi = [gx[1] + gx[end]; gx[2:(end-1)]]
- gxd[1] = (hx[1] + hx[end]) / 2
- gxd[end] = (hx[1] + hx[end]) / 2
- diagpos = 0
-
- if order4
- hxd3 = [hx3[end-1]; hx3[end]; hx3; hx3[1]]
- hxd13 = [hx[end-1]; hx[end]; hx; hx[1]]
- gxd3 = [gx3[end-1]; gx3; gx3[2]]
- gxd13 = [gx[end-1]; gx[1] + gx[end]; gx[2:(end-1)]; gx[end] + gx[1]; gx[2]]
- gxi3 = gx3[1:(end-1)]
- end
- end
-
- Bmap = spdiagm(Nux_in + 1, Nx + 2, diagpos => ones(T, Nux_in + 1))
-
- # Matrix to map from Nvx_t-1 to Nux_in points
- # (used in interpolation, convection_diffusion, viscosity)
- Bvux = spdiagm(Nux_in, Nvx_t - 1, diagpos => ones(T, Nux_in))
-
- # Map from Npx+2 points to Nux_t-1 points (ux faces)
- Bkux = copy(Bmap)
-
- ## Y-direction
-
- # gyi: integration and gyd: differentiation
- gyd = copy(gy)
- gyd[1] = hy[1]
- gyd[end] = hy[end]
-
- # hyi: integration and hyd: differentiation
- # Map to find suitable size
- hyi = copy(hy)
-
- # Restrict Ny+2 to Nvy_in+1 points
- if boundary_conditions.v.y == (:dirichlet, :dirichlet)
- yin = y[2:(end-1)]
- hyd = copy(hy)
- gyi = gy[2:(end-1)]
- diagpos = 1
-
- if order4
- hyd3 = [hy3[1]; hy3; hy3[end]]
- hyd13 = [hy[1]; hy; hy[end]]
- gyd3 = [2 * gy[1] + gy[2] + gy[3]; gy3; 2 * gy[end] + gy[end-1] + gy[end-2]]
- gyd13 = [gy[2]; 2 * gy[1]; gy[2:(end-1)]; 2 * gy[end]; gy[end-1]]
- gyi3 = gy3[2:(end-1)]
- end
- elseif boundary_conditions.v.y == (:dirichlet, :pressure)
- yin = y[2:end]
- hyd = [hy; hy[end]]
- gyi = gy[2:end]
- diagpos = 1
- elseif boundary_conditions.v.y == (:pressure, :dirichlet)
- yin = y[1:(end-1)]
- hyd = [hy[1]; hy]
- gyi = gy[1:(end-1)]
- diagpos = 0
- elseif boundary_conditions.v.y == (:pressure, :pressure)
- yin = y[1:end]
- hyd = [hy[1]; hy; hy[end]]
- gyi = copy(gy)
- diagpos = 0
- elseif boundary_conditions.v.y == (:periodic, :periodic)
- yin = y[1:(end-1)]
- hyd = [hy[end]; hy]
- gyi = [gy[1] + gy[end]; gy[2:(end-1)]]
- gyd[1] = (hy[1] + hy[end]) / 2
- gyd[end] = (hy[1] + hy[end]) / 2
- diagpos = 0
-
- if order4
- hyd3 = [hy3[end-1]; hy3[end]; hy3; hy3[1]]
- hyd13 = [hy[end-1]; hy[end]; hy; hy[1]]
- gyd3 = [gy3[end-1]; gy3; gy3[2]]
- gyd13 = [gy[end-1]; gy[1] + gy[end]; gy[2:(end-1)]; gy[end] + gy[1]; gy[2]]
- gyi3 = gy3[1:(end-1)]
- end
- end
-
- Bmap = spdiagm(Nvy_in + 1, Ny + 2, diagpos => ones(T, Nvy_in + 1))
-
- # Matrix to map from Nuy_t-1 to Nvy_in points
- # (used in interpolation, convection_diffusion)
- Buvy = spdiagm(Nvy_in, Nuy_t - 1, diagpos => ones(T, Nvy_in))
-
- # Map from Npy+2 points to Nvy_t-1 points (vy faces)
- Bkvy = copy(Bmap)
-
- ##
- # Volume (area) of pressure control volumes
- Ωp = hyi ⊗ hxi
-
- # Volume (area) of u control volumes
- Ωu = hyi ⊗ gxi
-
- # Volume of ux volumes
- Ωux = hyi ⊗ hxd
-
- # Volume of uy volumes
- Ωuy = gyd ⊗ gxi
-
- # Volume (area) of v control volumes
- Ωv = gyi ⊗ hxi
-
- # Volume of vx volumes
- Ωvx = gyi ⊗ gxd
-
- # Volume of vy volumes
- Ωvy = hyd ⊗ hxi
-
- Ω = [Ωu; Ωv]
-
- if order4
- # Differencing for second order operators on the fourth order mesh
- Ωux1 = hyi ⊗ hxd13
- Ωuy1 = gyd13 ⊗ gxi
- Ωvx1 = gyi ⊗ gxd13
- Ωvy1 = hyd13 ⊗ hxi
-
- # Volume (area) of pressure control volumes
- Ωp3 = hyi3 ⊗ hxi3
-
- # Volume (area) of u-vel control volumes
- Ωu3 = hyi3 ⊗ gxi3
-
- # Volume (area) of v-vel control volumes
- Ωv3 = gyi3 ⊗ hxi3
-
- # Volume (area) of dudx control volumes
- Ωux3 = hyi3 ⊗ hxd3
-
- # Volume (area) of dudy control volumes
- Ωuy3 = gyd3 ⊗ gxi3
-
- # Volume (area) of dvdx control volumes
- Ωvx3 = gyi3 ⊗ gxd3
-
- # Volume (area) of dvdy control volumes
- Ωvy3 = hyd3 ⊗ hxi3
-
- Ωu1 = copy(Ωu)
- Ωv1 = copy(Ωv)
-
- Ωu = α * Ωu1 - Ωu3
- Ωv = α * Ωv1 - Ωv3
- Ω = [Ωu; Ωv]
-
- Ωux = @. α * Ωux1 - Ωux3
- Ωuy = @. α * Ωuy1 - Ωuy3
- Ωvx = @. α * Ωvx1 - Ωvx3
- Ωvy = @. α * Ωvy1 - Ωvy3
- end
-
- # Metrics that can be useful for initialization:
- xu = ones(T, 1, Nuy_in) ⊗ xin
- yu = yp ⊗ ones(T, Nux_in)
- xu = reshape(xu, Nux_in, Nuy_in)
- yu = reshape(yu, Nux_in, Nuy_in)
-
- xv = ones(T, 1, Nvy_in) ⊗ xp
- yv = yin ⊗ ones(T, Nvx_in)
- xv = reshape(xv, Nvx_in, Nvy_in)
- yv = reshape(yv, Nvx_in, Nvy_in)
-
- xpp = ones(T, Ny) ⊗ xp
- ypp = yp ⊗ ones(T, Nx)
- xpp = reshape(xpp, Nx, Ny)
- ypp = reshape(ypp, Nx, Ny)
-
- # Indices of unknowns in velocity vector
- indu = 1:Nu
- indv = Nu .+ (1:Nv)
- indV = 1:NV
- indp = NV .+ (1:Np)
-
- ## Store quantities in the structure
- params = (;
- Npx,
- Npy,
- Np,
- Nux_in,
- Nux_b,
- Nux_t,
- Nuy_in,
- Nuy_b,
- Nuy_t,
- Nvx_in,
- Nvx_b,
- Nvx_t,
- Nvy_in,
- Nvy_b,
- Nvy_t,
- Nu,
- Nv,
- NV,
- Ωp,
- Ω,
- Ωux,
- Ωvx,
- Ωuy,
- Ωvy,
- hxi,
- hyi,
- hxd,
- hyd,
- gxi,
- gyi,
- gxd,
- gyd,
- Buvy,
- Bvux,
- Bkux,
- Bkvy,
- xin,
- yin,
- xu,
- yu,
- xv,
- yv,
- xpp,
- ypp,
- indu,
- indv,
- indV,
- indp,
- )
-
- if order4
- params = (;
- params...,
- hx3,
- hy3,
- hxi3,
- hyi3,
- gxi3,
- gyi3,
- hxd13,
- hxd3,
- hyd13,
- hyd3,
- gxd13,
- gxd3,
- gyd13,
- gyd3,
- Ωux1,
- Ωux3,
- Ωuy1,
- Ωuy3,
- Ωvx1,
- Ωvx3,
- Ωvy1,
- Ωvy3,
- )
- end
-
- (;
- dimension,
- α,
- β,
- order4,
- Nx,
- Ny,
- xlims,
- ylims,
- x,
- y,
- xp,
- yp,
- hx,
- hy,
- gx,
- gy,
- params...,
- )
-end
-
-"""
- Grid(x, y, z; boundary_conditions, order4 = false)
-
-Create nonuniform Cartesian box mesh `x` × `y` × `z` with boundary conditions `boundary_conditions`.
-If `order4` is `true`, a fourth order mesh is created.
-"""
-function Grid(x, y, z; boundary_conditions, order4 = false)
- dimension = Dimension(3)
-
- T = eltype(x)
-
- order4 && error("Fourth order grids not yet implemented for 3D")
- α = 81
- β = T(9 // 8)
-
- Nx = length(x) - 1
- Ny = length(y) - 1
- Nz = length(z) - 1
- xlims = (x[1], x[end])
- ylims = (y[1], y[end])
- zlims = (z[1], z[end])
-
- # Pressure positions
- xp = (x[1:(end-1)] + x[2:end]) / 2
- yp = (y[1:(end-1)] + y[2:end]) / 2
- zp = (z[1:(end-1)] + z[2:end]) / 2
-
- # Distance between velocity points
- hx = diff(x)
- hy = diff(y)
- hz = diff(z)
-
- # Distance between pressure points
- gx = zeros(T, Nx + 1)
- gx[1] = hx[1] / 2
- gx[2:Nx] = (hx[1:(Nx-1)] + hx[2:Nx]) / 2
- gx[Nx+1] = hx[end] / 2
-
- gy = zeros(T, Ny + 1)
- gy[1] = hy[1] / 2
- gy[2:Ny] = (hy[1:(Ny-1)] + hy[2:Ny]) / 2
- gy[Ny+1] = hy[end] / 2
-
- gz = zeros(T, Nz + 1)
- gz[1] = hz[1] / 2
- gz[2:Nz] = (hz[1:(Nz-1)] + hz[2:Nz]) / 2
- gz[Nz+1] = hz[end] / 2
-
- # Number of pressure points
- Npx = Nx
- Npy = Ny
- Npz = Nz
- Np = Npx * Npy * Npz
-
- ## u-volumes
- # x[1] x[2] x[3] .... x[Nx] x[Nx+1]
- # | | | | |
- # | | | | |
- # Dirichlet BC:
- # ULe u[1] u[2] .... u(Nx-1) uRi
- # Periodic BC:
- # u[1] u[2] u[3] .... u[Nx] u[1]
- # Pressure BC:
- # u[1] u[2] u[3] .... u[Nx] u[Nx+1]
-
- # x-dir
- Nux_b = 2 # Boundary points
- Nux_in = Nx + 1 # Inner points
- Nux_in -= boundary_conditions.u.x[1] ∈ [:dirichlet, :symmetric]
- Nux_in -= boundary_conditions.u.x[2] ∈ [:dirichlet, :symmetric]
- Nux_in -= boundary_conditions.u.x == (:periodic, :periodic)
- Nux_t = Nux_in + Nux_b # Total number
-
- # y-dir
- Nuy_b = 2 # Boundary points
- Nuy_in = Ny # Inner points
- Nuy_t = Nuy_in + Nuy_b # Total number
-
- # z-dir
- Nuz_b = 2 # Boundary points
- Nuz_in = Nz # Inner points
- Nuz_t = Nuz_in + Nuz_b # Total number
-
- # Total number
- Nu = Nux_in * Nuy_in * Nuz_in
-
- ## v-volumes
-
- # x-dir
- Nvx_b = 2 # Boundary points
- Nvx_in = Nx # Inner points
- Nvx_t = Nvx_in + Nvx_b # Total number
-
- # y-dir
- Nvy_b = 2 # Boundary points
- Nvy_in = Ny + 1 # Inner points
- Nvy_in -= boundary_conditions.v.y[1] ∈ [:dirichlet, :symmetric]
- Nvy_in -= boundary_conditions.v.y[2] ∈ [:dirichlet, :symmetric]
- Nvy_in -= boundary_conditions.v.y == (:periodic, :periodic)
- Nvy_t = Nvy_in + Nvy_b # Total number
-
- # z-dir
- Nvz_b = 2 # Boundary points
- Nvz_in = Nz # Inner points
- Nvz_t = Nvz_in + Nvz_b # Total number
-
- # Total number
- Nv = Nvx_in * Nvy_in * Nvz_in
-
- ## w-volumes
-
- # x-dir
- Nwx_b = 2 # Boundary points
- Nwx_in = Nx # Inner points
- Nwx_t = Nwx_in + Nwx_b # Total number
-
- # y-dir
- Nwy_b = 2 # Boundary points
- Nwy_in = Ny # Inner points
- Nwy_t = Nwy_in + Nwy_b # Total number
-
- # z-dir
- Nwz_b = 2 # Boundary points
- Nwz_in = Nz + 1 # Inner points
- Nwz_in -= boundary_conditions.w.z[1] ∈ [:dirichlet, :symmetric]
- Nwz_in -= boundary_conditions.w.z[2] ∈ [:dirichlet, :symmetric]
- Nwz_in -= boundary_conditions.w.z == (:periodic, :periodic)
- Nwz_t = Nwz_in + Nwz_b # Total number
-
- # Total number
- Nw = Nwx_in * Nwy_in * Nwz_in
-
- # Total number of velocity points
- NV = Nu + Nv + Nw
-
- ## Adapt mesh metrics depending on number of volumes
-
- ## X-direction
-
- # gxd: differentiation
- gxd = copy(gx)
- gxd[1] = hx[1]
- gxd[end] = hx[end]
-
- # hxi: integration and hxd: differentiation
- # Map to find suitable size
- hxi = copy(hx)
- hxd = [hx[1]; hx; hx[end]]
-
- # Restrict Nx+2 to Nux_in+1 points
- boundary_conditions.u.x == (:dirichlet, :dirichlet) && (diagpos = 1)
- boundary_conditions.u.x == (:dirichlet, :pressure) && (diagpos = 1)
- boundary_conditions.u.x == (:pressure, :dirichlet) && (diagpos = 0)
- boundary_conditions.u.x == (:pressure, :pressure) && (diagpos = 0)
- boundary_conditions.u.x == (:periodic, :periodic) && (diagpos = 0)
-
- Bmap = spdiagm(Nux_in + 1, Nx + 2, diagpos => ones(T, Nux_in + 1))
- Bmap_x_xin = spdiagm(Nux_in, Nx + 1, diagpos => ones(T, Nux_in))
- hxd = Bmap * hxd
- gxi = Bmap_x_xin * gx
- xin = Bmap_x_xin * x
-
- if boundary_conditions.u.x == (:periodic, :periodic)
- gxd[1] = (hx[1] + hx[end]) / 2
- gxd[end] = gxd[1]
- gxi[1] = gxd[1]
- hxd[1] = hx[end]
- end
-
- # Matrix to map from Nvx_t-1 to Nux_in points
- Bvux = spdiagm(Nux_in, Nvx_t - 1, diagpos => ones(T, Nux_in))
-
- # Matrix to map from Nwx_t-1 to Nux_in points
- Bwux = spdiagm(Nux_in, Nwx_t - 1, diagpos => ones(T, Nux_in))
-
- # Map from Npx+2 points to Nux_t-1 points (ux faces)
- Bkux = copy(Bmap)
-
- ## y-direction
-
- # gyi: integration and gyd: differentiation
- gyd = copy(gy)
- gyd[1] = hy[1]
- gyd[end] = hy[end]
-
- # hyi: integration and hyd: differentiation
- # Map to find suitable size
- hyi = copy(hy)
- hyd = [hy[1]; hy; hy[end]]
-
- # Restrict Ny+2 to Nvy_in+1 points
- boundary_conditions.v.y == (:dirichlet, :dirichlet) && (diagpos = 1)
- boundary_conditions.v.y == (:dirichlet, :pressure) && (diagpos = 1)
- boundary_conditions.v.y == (:pressure, :dirichlet) && (diagpos = 0)
- boundary_conditions.v.y == (:pressure, :pressure) && (diagpos = 0)
- boundary_conditions.v.y == (:periodic, :periodic) && (diagpos = 0)
-
- Bmap = spdiagm(Nvy_in + 1, Ny + 2, diagpos => ones(T, Nvy_in + 1))
- Bmap_y_yin = spdiagm(Nvy_in, Ny + 1, diagpos => ones(T, Nvy_in))
-
- hyd = Bmap * hyd
- gyi = Bmap_y_yin * gy
- yin = Bmap_y_yin * y
-
- if boundary_conditions.v.y == (:periodic, :periodic)
- gyd[1] = (hy[1] + hy[end]) / 2
- gyd[end] = gyd[1]
- gyi[1] = gyd[1]
- hyd[1] = hy[end]
- end
-
- # Matrix to map from Nuy_t-1 to Nvy_in points
- Buvy = spdiagm(Nvy_in, Nuy_t - 1, diagpos => ones(T, Nvy_in))
-
- # matrix to map from Nwy_t-1 to Nvy_in points
- Bwvy = spdiagm(Nvy_in, Nwy_t - 1, diagpos => ones(T, Nvy_in))
-
- # Map from Npy+2 points to Nvy_t-1 points (vy faces)
- Bkvy = copy(Bmap)
-
- ## z-direction
-
- # gzi: integration and gzd: differentiation
- gzd = copy(gz)
- gzd[1] = hz[1]
- gzd[end] = hz[end]
-
- # hzi: integration and hzd: differentiation
- # Map to find suitable size
- hzi = copy(hz)
- hzd = [hz[1]; hz; hz[end]]
-
- # Restrict Nz+2 to Nvz_in+1 points
- boundary_conditions.w.z == (:dirichlet, :dirichlet) && (diagpos = 1)
- boundary_conditions.w.z == (:dirichlet, :pressure) && (diagpos = 1)
- boundary_conditions.w.z == (:pressure, :dirichlet) && (diagpos = 0)
- boundary_conditions.w.z == (:pressure, :pressure) && (diagpos = 0)
- boundary_conditions.w.z == (:periodic, :periodic) && (diagpos = 0)
-
- shape = (Nwz_in + 1, Nz + 2)
- Bmap = spdiagm(shape..., diagpos => ones(T, Nwz_in + 1))
- Bmap_z_zin = spdiagm(Nwz_in, Nz + 1, diagpos => ones(T, Nwz_in))
- hzd = Bmap * hzd
- gzi = Bmap_z_zin * gz
- zin = Bmap_z_zin * z
-
- if boundary_conditions.w.z == (:periodic, :periodic)
- gzd[1] = (hz[1] + hz[end]) / 2
- gzd[end] = gzd[1]
- gzi[1] = gzd[1]
- hzd[1] = hz[end]
- end
-
- # Matrix to map from Nuz_t-1 to Nwz_in points
- Buwz = spdiagm(Nwz_in, Nuz_t - 1, diagpos => ones(T, Nwz_in))
-
- # Matrix to map from Nvz_t-1 to Nwz_in points
- Bvwz = spdiagm(Nwz_in, Nvz_t - 1, diagpos => ones(T, Nwz_in))
-
- # Map from Npy+2 points to Nvy_t-1 points (vy faces)
- Bkwz = copy(Bmap)
-
- ## Volumes
- # Volume (area) of pressure control volumes
- Ωp = hzi ⊗ hyi ⊗ hxi
-
- # Volume (area) of u control volumes
- Ωu = hzi ⊗ hyi ⊗ gxi
-
- # Volume (area) of v control volumes
- Ωv = hzi ⊗ gyi ⊗ hxi
-
- # Volume (area) of w control volumes
- Ωw = gzi ⊗ hyi ⊗ hxi
-
- # Total volumes
- Ω = [Ωu; Ωv; Ωw]
-
- # Metrics that can be useful for initialization:
- xu = reshape(ones(T, Nuz_in) ⊗ ones(T, Nuy_in) ⊗ xin, Nux_in, Nuy_in, Nuz_in)
- yu = reshape(ones(T, Nuz_in) ⊗ yp ⊗ ones(T, Nux_in), Nux_in, Nuy_in, Nuz_in)
- zu = reshape(zp ⊗ ones(T, Nuy_in) ⊗ ones(T, Nux_in), Nux_in, Nuy_in, Nuz_in)
-
- xv = reshape(ones(T, Nvz_in) ⊗ ones(T, Nvy_in) ⊗ xp, Nvx_in, Nvy_in, Nvz_in)
- yv = reshape(ones(T, Nvz_in) ⊗ yin ⊗ ones(T, Nvx_in), Nvx_in, Nvy_in, Nvz_in)
- zv = reshape(zp ⊗ ones(T, Nvy_in) ⊗ ones(T, Nvx_in), Nvx_in, Nvy_in, Nvz_in)
-
- xw = reshape(ones(T, Nwz_in) ⊗ ones(T, Nwy_in) ⊗ xp, Nwx_in, Nwy_in, Nwz_in)
- yw = reshape(ones(T, Nwz_in) ⊗ yp ⊗ ones(T, Nwx_in), Nwx_in, Nwy_in, Nwz_in)
- zw = reshape(zin ⊗ ones(T, Nwy_in) ⊗ ones(T, Nwx_in), Nwx_in, Nwy_in, Nwz_in)
-
- xpp = reshape(ones(T, Nz) ⊗ ones(T, Ny) ⊗ xp, Nx, Ny, Nz)
- ypp = reshape(ones(T, Nz) ⊗ yp ⊗ ones(T, Nx), Nx, Ny, Nz)
- zpp = reshape(zp ⊗ ones(T, Ny) ⊗ ones(T, Nx), Nx, Ny, Nz)
-
- # Indices of unknowns in velocity vector
- indu = 1:Nu
- indv = Nu .+ (1:Nv)
- indw = Nu + Nv .+ (1:Nw)
- indV = 1:NV
- indp = NV .+ (1:Np)
-
- (;
- dimension,
- order4,
- α,
- β,
- Nx,
- Ny,
- Nz,
- xlims,
- ylims,
- zlims,
- x,
- y,
- z,
- xp,
- yp,
- zp,
- hx,
- hy,
- hz,
- gx,
- gy,
- gz,
- Npx,
- Npy,
- Npz,
- Np,
- Nux_in,
- Nux_b,
- Nux_t,
- Nuy_in,
- Nuy_b,
- Nuy_t,
- Nuz_in,
- Nuz_b,
- Nuz_t,
- Nvx_in,
- Nvx_b,
- Nvx_t,
- Nvy_in,
- Nvy_b,
- Nvy_t,
- Nvz_in,
- Nvz_b,
- Nvz_t,
- Nwx_in,
- Nwx_b,
- Nwx_t,
- Nwy_in,
- Nwy_b,
- Nwy_t,
- Nwz_in,
- Nwz_b,
- Nwz_t,
- Nu,
- Nv,
- Nw,
- NV,
- Ωp,
- Ω,
- hxi,
- hyi,
- hzi,
- hxd,
- hyd,
- hzd,
- gxi,
- gyi,
- gzi,
- gxd,
- gyd,
- gzd,
- Buvy,
- Bvux,
- Buwz,
- Bwux,
- Bvwz,
- Bwvy,
- Bkux,
- Bkvy,
- Bkwz,
- xin,
- yin,
- zin,
- xu,
- yu,
- zu,
- xv,
- yv,
- zv,
- xw,
- yw,
- zw,
- xpp,
- ypp,
- zpp,
- indu,
- indv,
- indw,
- indV,
- indp,
- )
-end
diff --git a/src/grid/max_size.jl b/src/grid/max_size.jl
deleted file mode 100644
index 40961a1fa..000000000
--- a/src/grid/max_size.jl
+++ /dev/null
@@ -1,10 +0,0 @@
-"""
- max_size(grid)
-
-Get size of the largest grid element.
-"""
-function max_size(grid)
- (; hx, hy) = grid
- Δx, Δy = maximum(hx), maximum(hy)
- √(Δx^2 + Δy^2)
-end
diff --git a/src/grid/stretched_grid.jl b/src/grid/stretched_grid.jl
deleted file mode 100644
index 91ef76790..000000000
--- a/src/grid/stretched_grid.jl
+++ /dev/null
@@ -1,25 +0,0 @@
-"""
- stretched_grid(a, b, N, s = 1)
-
-Create a nonuniform grid of `N + 1` points from `a` to `b` with a stretch
-factor of `s`. If `s = 1`, return a uniform spacing from `a` to `b`. Otherwise,
-return a vector ``x \\in \\mathbb{R}^{N + 1}`` such that ``x_n = a + \\sum_{i =
-1}^n s^{i - 1} h`` for ``n = 0, \\dots , N``. Setting ``x_N = b`` then gives
-``h = (b - a) \\frac{1 - s}{1 - s^N}``, resulting in
-
-```math
-x_n = a + (b - a) \\frac{1 - s^n}{1 - s^N}, \\quad n = 0, \\dots, N.
-```
-
-Note that `stretched_grid(a, b, N, s)[n]` corresponds to ``x_{n - 1}``.
-
-See also [`cosine_grid`](@ref).
-"""
-function stretched_grid(a, b, N, s = 1)
- s > 0 || error("The strecth factor must be positive")
- if s ≈ 1
- LinRange(a, b, N + 1)
- else
- map(i -> a + (b - a) * (1 - s^i) / (1 - s^N), 0:N)
- end
-end
diff --git a/src/models/convection_models.jl b/src/models/convection_models.jl
deleted file mode 100644
index 10e47f340..000000000
--- a/src/models/convection_models.jl
+++ /dev/null
@@ -1,34 +0,0 @@
-"""
- AbstractConvectionModel
-
-Abstract convection model.
-"""
-abstract type AbstractConvectionModel end
-
-"""
- NoRegConvection()
-
-Unregularized convection model.
-"""
-struct NoRegConvectionModel <: AbstractConvectionModel end
-
-"""
- C2ConvectionModel()
-
-C2 regularization convection model.
-"""
-struct C2ConvectionModel <: AbstractConvectionModel end
-
-"""
- C4ConvectionModel()
-
-C4 regularization convection model.
-"""
-struct C4ConvectionModel <: AbstractConvectionModel end
-
-"""
- LerayConvectionModel()
-
-Leray regularization convection model.
-"""
-struct LerayConvectionModel <: AbstractConvectionModel end
diff --git a/src/models/viscosity_models.jl b/src/models/viscosity_models.jl
deleted file mode 100644
index f5dd24282..000000000
--- a/src/models/viscosity_models.jl
+++ /dev/null
@@ -1,44 +0,0 @@
-"""
- AbstractViscosityModel
-
-Abstract viscosity model.
-"""
-abstract type AbstractViscosityModel{T} end
-
-"""
- LaminarModel(Re)
-
-Laminar model with Reynolds number `Re`.
-"""
-Base.@kwdef struct LaminarModel{T} <: AbstractViscosityModel{T}
- Re::T # Reynolds number
-end
-
-"""
- MixingLengthModel(Re)
-
-Mixing-length model with Reynolds number `Re` and mixing length `lm`.
-"""
-Base.@kwdef struct MixingLengthModel{T} <: AbstractViscosityModel{T}
- Re::T # Reynolds number
- lm::T = 1 # Mixing length
-end
-
-"""
- SmagorinskyModel(Re, C_s = 0.17)
-
-Smagorinsky-Lilly model with Reynolds number `Re` and constant `C_s`.
-"""
-Base.@kwdef struct SmagorinskyModel{T} <: AbstractViscosityModel{T}
- Re::T # Reynolds number
- C_s::T = 0.17 # Smagorinsky constant
-end
-
-"""
- QR(Re)
-
-QR-model with Reynolds number `Re`.
-"""
-Base.@kwdef struct QRModel{T} <: AbstractViscosityModel{T}
- Re::T # Reynolds number
-end
diff --git a/src/momentum/check_symmetry.jl b/src/momentum/check_symmetry.jl
deleted file mode 100644
index 44e5def85..000000000
--- a/src/momentum/check_symmetry.jl
+++ /dev/null
@@ -1,45 +0,0 @@
-"""
- check_symmetry(V, t, setup, ϵ = 1e-14)
-
-Check symmetry of convection operator.
-
-`flag = 0`: no symmetry error
-`flag = 1`: symmetry error
-"""
-function check_symmetry(V, t, setup, ϵ = 1e-14; bc_vectors)
- (; grid, operators, boundary_conditions) = setup
- (; indu, indv) = grid
- (; Cux, Cuy, Cvx, Cvy) = operators
- (; Au_ux, Au_uy, Av_vx, Av_vy) = operators
- (; Iu_ux, Iv_uy, Iu_vx, Iv_vy) = operators
- (; yIu_ux, yIv_uy, yIu_vx, yIv_vy) = bc_vectors
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- Cu =
- Cux * spdiagm(Iu_ux * uₕ + yIu_ux) * Au_ux +
- Cuy * spdiagm(Iv_uy * vₕ + yIv_uy) * Au_uy
- Cv =
- Cvx * spdiagm(Iu_vx * uₕ + yIu_vx) * Av_vx +
- Cvy * spdiagm(Iv_vy * vₕ + yIv_vy) * Av_vy
-
- error_u = maximum(abs.(Cu + Cu'))
- error_v = maximum(abs.(Cv + Cv'))
-
- symmetry_error = max(error_u, error_v)
-
- flag = 0
- if symmetry_error > ϵ
- if boundary_conditions.u.x[1] != :pressure &&
- boundary_conditions.u.x[2] != :pressure
- flag = 1
- end
- if boundary_conditions.v.y[1] != :pressure &&
- boundary_conditions.v.y[2] != :pressure
- flag = 1
- end
- end
-
- flag, symmetry_error
-end
diff --git a/src/momentum/compute_conservation.jl b/src/momentum/compute_conservation.jl
deleted file mode 100644
index 5b8c7757a..000000000
--- a/src/momentum/compute_conservation.jl
+++ /dev/null
@@ -1,124 +0,0 @@
-"""
- compute_conservation(V, t, setup; bc_vectors = nothing)
-
-Compute mass, momentum and energy conservation properties of velocity field.
-"""
-function compute_conservation end
-
-compute_conservation(V, t, setup; bc_vectors = nothing) =
- compute_conservation(setup.grid.dimension, V, t, setup; bc_vectors = nothing)
-
-# 2D version
-function compute_conservation(::Dimension{2}, V, t, setup; bc_vectors = nothing)
- (; grid, operators, boundary_conditions) = setup
- (; indu, indv, Ω, x, y, xp, yp, hx, hy, gx, gy) = grid
- (; M) = operators
- (; bc_unsteady, u_bc, v_bc) = boundary_conditions
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- Ωu = @view Ω[indu]
- Ωv = @view Ω[indv]
-
- if isnothing(bc_vectors) || bc_unsteady
- bc_vectors = get_bc_vectors(setup, t)
- end
- (; yM) = bc_vectors
-
- uLe_i = reshape(u_bc.(x[1], yp, t), :)
- uRi_i = reshape(u_bc.(x[end], yp, t), :)
- vLo_i = reshape(v_bc.(xp, y[1], t), :)
- vUp_i = reshape(v_bc.(xp, y[end], t), :)
-
- # Check if new velocity field is divergence free (mass conservation)
- maxdiv = maximum(abs.(M * V + yM))
-
- # Calculate total momentum
- umom = sum(Ωu .* uₕ)
- vmom = sum(Ωv .* vₕ)
-
- # Add boundary contributions in case of Dirichlet BC
- boundary_conditions.u.x[1] == :dirichlet && (umom += sum(uLe_i .* hy) * gx[1])
- boundary_conditions.u.x[2] == :dirichlet && (umom += sum(uRi_i .* hy) * gx[end])
- boundary_conditions.v.y[1] == :dirichlet && (vmom += sum(vLo_i .* hx) * gy[1])
- boundary_conditions.v.y[2] == :dirichlet && (vmom += sum(vUp_i .* hx) * gy[end])
-
- # Calculate total kinetic energy
- k = 1 // 2 * sum(Ω .* V .^ 2)
-
- # Add boundary contributions in case of Dirichlet BC
- boundary_conditions.u.x[1] == :dirichlet &&
- (k += 1 // 2 * sum(uLe_i .^ 2 .* hy) * gx[1])
- boundary_conditions.u.x[2] == :dirichlet &&
- (k += 1 // 2 * sum(uRi_i .^ 2 .* hy) * gx[end])
- boundary_conditions.v.y[1] == :dirichlet &&
- (k += 1 // 2 * sum(vLo_i .^ 2 .* hx) * gy[1])
- boundary_conditions.v.y[2] == :dirichlet &&
- (k += 1 // 2 * sum(vUp_i .^ 2 .* hx) * gy[end])
-
- maxdiv, umom, vmom, k
-end
-
-# 3D version
-function compute_conservation(::Dimension{3}, V, t, setup; bc_vectors = nothing)
- (; grid, operators, boundary_conditions) = setup
- (; indu, indv, indw, Ω, x, y, z, xp, yp, zp, hx, hy, hz, gx, gy, gz) = grid
- (; M) = operators
- (; bc_unsteady, u_bc, v_bc, w_bc) = boundary_conditions
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- wₕ = @view V[indw]
-
- Ωu = @view Ω[indu]
- Ωv = @view Ω[indv]
- Ωw = @view Ω[indw]
-
- if isnothing(bc_vectors) || bc_unsteady
- bc_vectors = get_bc_vectors(setup, t)
- end
- (; yM) = bc_vectors
-
- uLe_i = reshape(u_bc.(x[1], yp, zp', t), :)
- uRi_i = reshape(u_bc.(x[end], yp, zp', t), :)
- vLo_i = reshape(v_bc.(xp, y[1], zp', t), :)
- vUp_i = reshape(v_bc.(xp, y[end], zp', t), :)
- wBa_i = reshape(w_bc.(xp, yp', z[1], t), :)
- wFr_i = reshape(w_bc.(xp, yp', z[end], t), :)
-
- # Check if new velocity field is divergence free (mass conservation)
- maxdiv = maximum(abs.(M * V + yM))
-
- # Calculate total momentum
- umom = sum(Ωu .* uₕ)
- vmom = sum(Ωv .* vₕ)
- wmom = sum(Ωw .* wₕ)
-
- # Add boundary contributions in case of Dirichlet BC
- boundary_conditions.u.x[1] == :dirichlet && (umom += sum(uLe_i .* (hz ⊗ hy)) * gx[1])
- boundary_conditions.u.x[2] == :dirichlet && (umom += sum(uRi_i .* (hz ⊗ hy)) * gx[end])
- boundary_conditions.v.y[1] == :dirichlet && (vmom += sum(vLo_i .* (hz ⊗ hx)) * gy[1])
- boundary_conditions.v.y[2] == :dirichlet && (vmom += sum(vUp_i .* (hz ⊗ hx)) * gy[end])
- boundary_conditions.w.z[1] == :dirichlet && (wmom += sum(wBa_i .* (hy ⊗ hx)) * gz[1])
- boundary_conditions.w.z[2] == :dirichlet && (wmom += sum(wFr_i .* (hy ⊗ hx)) * gz[end])
-
- # Calculate total kinetic energy
- k = 1 // 2 * sum(Ω .* V .^ 2)
-
- # Add boundary contributions in case of Dirichlet BC
- boundary_conditions.u.x[1] == :dirichlet &&
- (k += 1 // 2 * sum(uLe_i .^ 2 .* (hz ⊗ hy)) * gx[1])
- boundary_conditions.u.x[2] == :dirichlet &&
- (k += 1 // 2 * sum(uRi_i .^ 2 .* (hz ⊗ hy)) * gx[end])
- boundary_conditions.v.y[1] == :dirichlet &&
- (k += 1 // 2 * sum(vLo_i .^ 2 .* (hz ⊗ hx)) * gy[1])
- boundary_conditions.v.y[2] == :dirichlet &&
- (k += 1 // 2 * sum(vUp_i .^ 2 .* (hz ⊗ hx)) * gy[end])
- boundary_conditions.w.z[1] == :dirichlet &&
- (k += 1 // 2 * sum(wBa_i .^ 2 .* (hy ⊗ hx)) * gz[1])
- boundary_conditions.w.z[2] == :dirichlet &&
- (k += 1 // 2 * sum(wFr_i .^ 2 .* (hy ⊗ hx)) * gz[end])
-
- maxdiv, umom, vmom, wmom, k
-end
diff --git a/src/momentum/convection.jl b/src/momentum/convection.jl
deleted file mode 100644
index bcc33544f..000000000
--- a/src/momentum/convection.jl
+++ /dev/null
@@ -1,860 +0,0 @@
-"""
- convection(
- model, V, ϕ, setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
- )
-
-Evaluate convective terms `c` and, optionally, Jacobian `∇c = ∂c/∂V`, using the convection
-model `model`. The convected quantity is `ϕ` (usually `ϕ = V`).
-
-Non-mutating/allocating/out-of-place version.
-
-See also [`convection!`](@ref).
-"""
-function convection end
-
-convection(m, V, ϕ, setup; kwargs...) =
- convection(setup.grid.dimension, m, V, ϕ, setup; kwargs...)
-
-function convection(
- dimension,
- ::NoRegConvectionModel,
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; order4, α) = setup.grid
-
- # No regularization
- c, ∇c = convection_components(
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian,
- newton_factor,
- order4 = false,
- )
-
- if order4
- c3, ∇c3 = convection_components(
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian,
- newton_factor,
- order4,
- )
- c = @. α * c - c3
- get_jacobian && (∇c = @. α * ∇c - ∇c3)
- end
-
- c, ∇c
-end
-
-# 2D version
-function convection(
- ::Dimension{2},
- ::C2ConvectionModel,
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv) = grid
- (; Diffu_f, Diffv_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yM) = bc_vectors
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
-
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
-
- ūₕ = filter_convection(uₕ, Diffu_f, yDiffu_f, α_reg)
- v̄ₕ = filter_convection(vₕ, Diffv_f, yDiffv_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v]
- V̄ = [ūₕ; v̄ₕ]
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * ϕ̄ + yM))
-
- c, ∇c = convection_components(V̄, ϕ̄, setup; bc_vectors, get_jacobian, newton_factor)
-
- cu = @view c[indu]
- cv = @view c[indv]
-
- cu = filter_convection(cu, Diffu_f, yDiffu_f, α_reg)
- cv = filter_convection(cv, Diffv_f, yDiffv_f, α_reg)
-
- c = [cu; cv]
-
- c, ∇c
-end
-
-# 3D version
-function convection(
- ::Dimension{3},
- ::C2ConvectionModel,
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv, indw) = grid
- (; Diffu_f, Diffv_f, Diffw_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yDiffw_f, yM) = bc_vectors
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- wₕ = @view V[indw]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
- ϕw = @view ϕ[indw]
-
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
- ϕ̄w = filter_convection(ϕw, Diffw_f, yDiffw_f, α_reg)
-
- ūₕ = filter_convection(uₕ, Diffu_f, yDiffu_f, α_reg)
- v̄ₕ = filter_convection(vₕ, Diffv_f, yDiffv_f, α_reg)
- w̄ₕ = filter_convection(wₕ, Diffw_f, yDiffw_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v; ϕ̄w]
- V̄ = [ūₕ; v̄ₕ; w̄ₕ]
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * ϕ̄ + yM))
-
- c, ∇c =
- convection_components(V̄, ϕ̄, setup, cache; bc_vectors, get_jacobian, newton_factor)
-
- cu = @view c[indu]
- cv = @view c[indv]
- cw = @view c[indw]
-
- cu = filter_convection(cu, Diffu_f, yDiffu_f, α_reg)
- cv = filter_convection(cv, Diffv_f, yDiffv_f, α_reg)
- cw = filter_convection(cw, Diffw_f, yDiffw_f, α_reg)
-
- c = [cu; cv; cw]
-
- c, ∇c
-end
-
-# 2D version
-function convection(
- ::Dimension{2},
- ::C4ConvectionModel,
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv) = grid
- (; Diffu_f, Diffv_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yM) = bc_vectors
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
-
- # C4 consists of 3 terms:
- # C4 = conv(filter(u), filter(u)) + filter(conv(filter(u), u') +
- # filter(conv(u', filter(u)))
- # where u' = u - filter(u)
-
- # Filter both convecting and convected velocity
- ūₕ = filter_convection(uₕ, Diffu_f, yDiffu_f, α_reg)
- v̄ₕ = filter_convection(vₕ, Diffv_f, yDiffv_f, α_reg)
-
- V̄ = [ūₕ; v̄ₕ]
- ΔV = V - V̄
-
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v]
- Δϕ = ϕ - ϕ̄
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * V̄ + yM))
-
- c, ∇c = convection_components(V̄, ϕ̄, setup; bc_vectors, get_jacobian, newton_factor)
- c2, ∇c2 = convection_components(ΔV, ϕ̄, setup; bc_vectors, get_jacobian, newton_factor)
- c3, ∇c3 = convection_components(V̄, Δϕ, setup; bc_vectors, get_jacobian, newton_factor)
-
- cu = @view c[indu]
- cv = @view c[indv]
-
- cu2 = @view c2[indu]
- cv2 = @view c2[indv]
-
- cu3 = @view c3[indu]
- cv3 = @view c3[indv]
-
- cu += filter_convection(cu2 + cu3, Diffu_f, yDiffu_f, α_reg)
- cv += filter_convection(cv2 + cv3, Diffv_f, yDiffv_f, α_reg)
-
- c = [cu; cv]
-
- c, ∇c
-end
-
-# 3D version
-function convection(
- ::Dimension{3},
- ::C4ConvectionModel,
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv, indw) = grid
- (; Diffu_f, Diffv_f, Diffw_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yDiffw_f, yM) = bc_vectors
- (; c2, ∇c2, c3, ∇c3) = cache
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- wₕ = @view V[indw]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
- ϕw = @view ϕ[indw]
-
- # C4 consists of 3 terms:
- # C4 = conv(filter(u), filter(u)) + filter(conv(filter(u), u') +
- # filter(conv(u', filter(u)))
- # where u' = u - filter(u)
-
- # Filter both convecting and convected velocity
- ūₕ = filter_convection(uₕ, Diffu_f, yDiffu_f, α_reg)
- v̄ₕ = filter_convection(vₕ, Diffv_f, yDiffv_f, α_reg)
- w̄ₕ = filter_convection(wₕ, Diffw_f, yDiffw_f, α_reg)
-
- V̄ = [ūₕ; v̄ₕ; w̄ₕ]
- ΔV = V - V̄
-
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
- ϕ̄w = filter_convection(ϕw, Diffw_f, yDiffw_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v; ϕ̄w]
- Δϕ = ϕ - ϕ̄
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * V̄ + yM))
-
- c, ∇c = convection_components(V̄, ϕ̄, setup; bc_vectors, get_jacobian, newton_factor)
- c2, ∇c2 = convection_components(ΔV, ϕ̄, setup; bc_vectors, get_jacobian, newton_factor)
- c3, ∇c3 = convection_components(V̄, Δϕ, setup; bc_vectors, get_jacobian, newton_factor)
-
- cu = @view c[indu]
- cv = @view c[indv]
- cw = @view c[indw]
-
- cu2 = @view c2[indu]
- cv2 = @view c2[indv]
- cw2 = @view c2[indw]
-
- cu3 = @view c3[indu]
- cv3 = @view c3[indv]
- cw3 = @view c3[indw]
-
- cu += filter_convection(cu2 + cu3, Diffu_f, yDiffu_f, α_reg)
- cv += filter_convection(cv2 + cv3, Diffv_f, yDiffv_f, α_reg)
- cw += filter_convection(cw2 + cw3, Diffw_f, yDiffw_f, α_reg)
-
- c = [cu; cv; cw]
-
- c, ∇c
-end
-
-# 2D version
-function convection(
- ::Dimension{2},
- ::LerayConvectionModel,
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv) = grid
- (; Diffu_f, Diffv_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yM) = bc_vectors
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
-
- # TODO: needs finishing
- error("Leray convection not implemented yet")
-
- # Filter the convecting field
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v]
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * ϕ̄ + yM))
-
- convection_components(V, ϕ̄, setup; bc_vectors, get_jacobian, newton_factor)
-end
-
-# 3D version
-function convection(
- ::Dimension{3},
- ::LerayConvectionModel,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- bc_vectors,
- cache;
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv, indw) = grid
- (; Diffu_f, Diffv_f, Diffw_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yDiffw_f, yM) = bc_vectors
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
- ϕw = @view ϕ[indw]
-
- # TODO: needs finishing
- error("Leray convection not implemented yet")
-
- # Filter the convecting field
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
- ϕ̄w = filter_convection(ϕw, Diffw_f, yDiffw_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v; ϕ̄w]
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * ϕ̄ + yM))
-
- convection_components(V, ϕ̄, setup; bc_vectors, get_jacobian, newton_factor)
-end
-
-"""
- convection!(
- model, c, ∇c, V, ϕ, setup, cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
- )
-
-Evaluate convective terms `c` and, optionally, Jacobian `∇c = ∂c/∂V`, using the convection
-model `model`. The convected quantity is `ϕ` (usually `ϕ = V`).
-
-Mutating/non-allocating/in-place version.
-
-See also [`convection`](@ref).
-"""
-function convection! end
-
-convection!(m, c, ∇c, V, ϕ, setup, cache; kwargs...) =
- convection!(setup.grid.dimension, m, c, ∇c, V, ϕ, setup, cache; kwargs...)
-
-function convection!(
- dimension,
- ::NoRegConvectionModel,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; order4, α) = setup.grid
- (; c3, ∇c3) = cache
-
- # No regularization
- convection_components!(
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- order4 = false,
- )
-
- if order4
- convection_components!(
- c3,
- ∇c3,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- order4,
- )
- @. c = α * c - c3
- get_jacobian && (@. ∇c = α * ∇c - ∇c3)
- end
-
- c, ∇c
-end
-
-# 2D version
-function convection!(
- ::Dimension{2},
- ::C2ConvectionModel,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv) = grid
- (; Diffu_f, Diffv_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yM) = bc_vectors
-
- cu = @view c[indu]
- cv = @view c[indv]
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
-
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
-
- ūₕ = filter_convection(uₕ, Diffu_f, yDiffu_f, α_reg)
- v̄ₕ = filter_convection(vₕ, Diffv_f, yDiffv_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v]
- V̄ = [ūₕ; v̄ₕ]
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * ϕ̄ + yM))
-
- convection_components!(
- c,
- ∇c,
- V̄,
- ϕ̄,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
-
- cu .= filter_convection(cu, Diffu_f, yDiffu_f, α_reg)
- cv .= filter_convection(cv, Diffv_f, yDiffv_f, α_reg)
-
- c, ∇c
-end
-
-# 3D version
-function convection!(
- ::Dimension{3},
- ::C2ConvectionModel,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv, indw) = grid
- (; Diffu_f, Diffv_f, Diffw_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yDiffw_f, yM) = bc_vectors
-
- cu = @view c[indu]
- cv = @view c[indv]
- cw = @view c[indw]
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- wₕ = @view V[indw]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
- ϕw = @view ϕ[indw]
-
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
- ϕ̄w = filter_convection(ϕw, Diffw_f, yDiffw_f, α_reg)
-
- ūₕ = filter_convection(uₕ, Diffu_f, yDiffu_f, α_reg)
- v̄ₕ = filter_convection(vₕ, Diffv_f, yDiffv_f, α_reg)
- w̄ₕ = filter_convection(wₕ, Diffw_f, yDiffw_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v; ϕ̄w]
- V̄ = [ūₕ; v̄ₕ; w̄ₕ]
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * ϕ̄ + yM))
-
- convection_components!(
- c,
- ∇c,
- V̄,
- ϕ̄,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
-
- cu .= filter_convection(cu, Diffu_f, yDiffu_f, α_reg)
- cv .= filter_convection(cv, Diffv_f, yDiffv_f, α_reg)
- cw .= filter_convection(cw, Diffw_f, yDiffw_f, α_reg)
-
- c, ∇c
-end
-
-# 2D version
-function convection!(
- ::Dimension{2},
- ::C4ConvectionModel,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv) = grid
- (; Diffu_f, Diffv_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yM) = bc_vectors
- (; c2, ∇c2, c3, ∇c3) = cache
-
- cu = @view c[indu]
- cv = @view c[indv]
- cu2 = @view c2[indu]
- cv2 = @view c2[indv]
- cu3 = @view c3[indu]
- cv3 = @view c3[indv]
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
-
- # C4 consists of 3 terms:
- # C4 = conv(filter(u), filter(u)) + filter(conv(filter(u), u') +
- # filter(conv(u', filter(u)))
- # where u' = u - filter(u)
-
- # Filter both convecting and convected velocity
- ūₕ = filter_convection(uₕ, Diffu_f, yDiffu_f, α_reg)
- v̄ₕ = filter_convection(vₕ, Diffv_f, yDiffv_f, α_reg)
-
- V̄ = [ūₕ; v̄ₕ]
- ΔV = V - V̄
-
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v]
- Δϕ = ϕ - ϕ̄
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * V̄ + yM))
-
- convection_components!(
- c,
- ∇c,
- V̄,
- ϕ̄,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
- convection_components!(
- c2,
- ∇c2,
- ΔV,
- ϕ̄,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
- convection_components!(
- c3,
- ∇c3,
- V̄,
- Δϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
-
- cu .+= filter_convection(cu2 + cu3, Diffu_f, yDiffu_f, α_reg)
- cv .+= filter_convection(cv2 + cv3, Diffv_f, yDiffv_f, α_reg)
-
- c, ∇c
-end
-
-# 3D version
-function convection!(
- ::Dimension{3},
- ::C4ConvectionModel,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv, indw) = grid
- (; Diffu_f, Diffv_f, Diffw_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yDiffw_f, yM) = bc_vectors
- (; c2, ∇c2, c3, ∇c3) = cache
-
- cu = @view c[indu]
- cv = @view c[indv]
- cw = @view c[indw]
-
- cu2 = @view c2[indu]
- cv2 = @view c2[indv]
- cw2 = @view c2[indw]
-
- cu3 = @view c3[indu]
- cv3 = @view c3[indv]
- cw3 = @view c3[indw]
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- wₕ = @view V[indw]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
- ϕw = @view ϕ[indw]
-
- # C4 consists of 3 terms:
- # C4 = conv(filter(u), filter(u)) + filter(conv(filter(u), u') +
- # filter(conv(u', filter(u)))
- # where u' = u - filter(u)
-
- # Filter both convecting and convected velocity
- ūₕ = filter_convection(uₕ, Diffu_f, yDiffu_f, α_reg)
- v̄ₕ = filter_convection(vₕ, Diffv_f, yDiffv_f, α_reg)
- w̄ₕ = filter_convection(wₕ, Diffw_f, yDiffw_f, α_reg)
-
- V̄ = [ūₕ; v̄ₕ; w̄ₕ]
- ΔV = V - V̄
-
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
- ϕ̄w = filter_convection(ϕw, Diffw_f, yDiffw_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v; ϕ̄w]
- Δϕ = ϕ - ϕ̄
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * V̄ + yM))
-
- convection_components!(
- c,
- ∇c,
- V̄,
- ϕ̄,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
- convection_components!(
- c2,
- ∇c2,
- ΔV,
- ϕ̄,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
- convection_components!(
- c3,
- ∇c3,
- V̄,
- Δϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
-
- cu .+= filter_convection(cu2 + cu3, Diffu_f, yDiffu_f, α_reg)
- cv .+= filter_convection(cv2 + cv3, Diffv_f, yDiffv_f, α_reg)
- cw .+= filter_convection(cw2 + cw3, Diffw_f, yDiffw_f, α_reg)
-
- c, ∇c
-end
-
-# 2D version
-function convection!(
- ::Dimension{2},
- ::LerayConvectionModel,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv) = grid
- (; Diffu_f, Diffv_f, M, α_reg) = operators
- (; yDiffu_f, yDiffv_f, yM) = bc_vectors
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
-
- # TODO: needs finishing
- error("Leray convection not implemented yet")
-
- # Filter the convecting field
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α_reg)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α_reg)
-
- ϕ̄ = [ϕ̄u; ϕ̄v]
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * ϕ̄ + yM))
-
- convection_components!(
- c,
- ∇c,
- V,
- ϕ̄,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
-
- c, ∇c
-end
-
-# 3D version
-function convection!(
- ::Dimension{3},
- ::LerayConvectionModel,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
-)
- (; grid, operators) = setup
- (; indu, indv, indw) = grid
- (; Diffu_f, Diffv_f, Diffw_f, M, α) = operators
- (; yDiffu_f, yDiffv_f, yDiffw_f, yM) = bc_vectors
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
- ϕw = @view ϕ[indw]
-
- # TODO: needs finishing
- error("Leray convection not implemented yet")
-
- # Filter the convecting field
- ϕ̄u = filter_convection(ϕu, Diffu_f, yDiffu_f, α)
- ϕ̄v = filter_convection(ϕv, Diffv_f, yDiffv_f, α)
- ϕ̄w = filter_convection(ϕw, Diffw_f, yDiffw_f, α)
-
- ϕ̄ = [ϕ̄u; ϕ̄v; ϕ̄w]
-
- # Divergence of filtered velocity field; should be zero!
- maxdiv_f = maximum(abs.(M * ϕ̄ + yM))
-
- convection_components!(
- c,
- ∇c,
- V,
- ϕ̄,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
-
- c, ∇c
-end
diff --git a/src/momentum/convection_components.jl b/src/momentum/convection_components.jl
deleted file mode 100644
index 6783b7a3b..000000000
--- a/src/momentum/convection_components.jl
+++ /dev/null
@@ -1,634 +0,0 @@
-"""
- convection_components(
- V, ϕ, setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
- order4 = false,
- )
-
-Compute convection components.
-
-Non-mutating/allocating/out-of-place version.
-
-See also [`convection_components!`](@ref).
-"""
-function convection_components end
-
-convection_components(V, ϕ, setup; kwargs...) =
- convection_components(setup.grid.dimension, V, ϕ, setup; kwargs...)
-
-# 2D version
-function convection_components(
- ::Dimension{2},
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
- order4 = false,
-)
- (; grid, operators) = setup
-
- if order4
- Cux = operators.Cux3
- Cuy = operators.Cuy3
- Cvx = operators.Cvx3
- Cvy = operators.Cvy3
-
- Au_ux = operators.Au_ux3
- Au_uy = operators.Au_uy3
- Av_vx = operators.Av_vx3
- Av_vy = operators.Av_vy3
-
- Iu_ux = operators.Iu_ux3
- Iv_uy = operators.Iv_uy3
- Iu_vx = operators.Iu_vx3
- Iv_vy = operators.Iv_vy3
-
- yAu_ux = bc_vectors.yAu_ux3
- yAu_uy = bc_vectors.yAu_uy3
- yAv_vx = bc_vectors.yAv_vx3
- yAv_vy = bc_vectors.yAv_vy3
-
- yIu_ux = bc_vectors.yIu_ux3
- yIv_uy = bc_vectors.yIv_uy3
- yIu_vx = bc_vectors.yIu_vx3
- yIv_vy = bc_vectors.yIv_vy3
- else
- (; Cux, Cuy, Cvx, Cvy) = operators
- (; Au_ux, Au_uy, Av_vx, Av_vy) = operators
- (; Iu_ux, Iv_uy, Iu_vx, Iv_vy) = operators
- (; yAu_ux, yAu_uy, yAv_vx, yAv_vy) = bc_vectors
- (; yIu_ux, yIv_uy, yIu_vx, yIv_vy) = bc_vectors
- end
-
- (; indu, indv) = grid
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
-
- # Convection components
- u_ux = Au_ux * uₕ + yAu_ux # u at ux
- ū_ux = Iu_ux * ϕu + yIu_ux # ū at ux
- ∂uū∂x = Cux * (u_ux .* ū_ux)
-
- u_uy = Au_uy * uₕ + yAu_uy # u at uy
- v̄_uy = Iv_uy * ϕv + yIv_uy # ū at uy
- ∂uv̄∂y = Cuy * (u_uy .* v̄_uy)
-
- v_vx = Av_vx * vₕ + yAv_vx # v at vx
- ū_vx = Iu_vx * ϕu + yIu_vx # ū at vx
- ∂vū∂x = Cvx * (v_vx .* ū_vx)
-
- v_vy = Av_vy * vₕ + yAv_vy # v at vy
- v̄_vy = Iv_vy * ϕv + yIv_vy # ū at vy
- ∂vv̄∂y = Cvy * (v_vy .* v̄_vy)
-
- cu = ∂uū∂x + ∂uv̄∂y
- cv = ∂vū∂x + ∂vv̄∂y
-
- c = [cu; cv]
-
- if get_jacobian
- ## Convective terms, u-component
- C1 = Cux * Diagonal(ū_ux)
- C2 = Cux * Diagonal(u_ux) * newton_factor
- Conv_ux_11 = C1 * Au_ux .+ C2 * Iu_ux
-
- C1 = Cuy * Diagonal(v̄_uy)
- C2 = Cuy * Diagonal(u_uy) * newton_factor
- Conv_uy_11 = C1 * Au_uy
- Conv_uy_12 = C2 * Iv_uy
-
- ## Convective terms, v-component
- C1 = Cvx * Diagonal(ū_vx)
- C2 = Cvx * Diagonal(v_vx) * newton_factor
- Conv_vx_21 = C2 * Iu_vx
- Conv_vx_22 = C1 * Av_vx
-
- C1 = Cvy * Diagonal(v̄_vy)
- C2 = Cvy * Diagonal(v_vy) * newton_factor
- Conv_vy_22 = C1 * Av_vy .+ C2 * Iv_vy
-
- ∇c = [
- (Conv_ux_11+Conv_uy_11) Conv_uy_12
- Conv_vx_21 (Conv_vx_22+Conv_vy_22)
- ]
- else
- ∇c = nothing
- end
-
- c, ∇c
-end
-
-# 3D version
-function convection_components(
- ::Dimension{3},
- V,
- ϕ,
- setup;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
- order4 = false,
-)
- order4 && error("order4 not implemented for 3D")
-
- (; grid, operators) = setup
- (; indu, indv, indw) = grid
- (; Cux, Cuy, Cuz, Cvx, Cvy, Cvz, Cwx, Cwy, Cwz) = operators
-
- (; Au_ux, Au_uy, Au_uz) = operators
- (; Av_vx, Av_vy, Av_vz) = operators
- (; Aw_wx, Aw_wy, Aw_wz) = operators
-
- (; Iu_ux, Iv_uy, Iw_uz) = operators
- (; Iu_vx, Iv_vy, Iw_vz) = operators
- (; Iu_wx, Iv_wy, Iw_wz) = operators
-
- (; yAu_ux, yAu_uy, yAu_uz) = bc_vectors
- (; yAv_vx, yAv_vy, yAv_vz) = bc_vectors
- (; yAw_wx, yAw_wy, yAw_wz) = bc_vectors
-
- (; yIu_ux, yIv_uy, yIw_uz) = bc_vectors
- (; yIu_vx, yIv_vy, yIw_vz) = bc_vectors
- (; yIu_wx, yIv_wy, yIw_wz) = bc_vectors
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- wₕ = @view V[indw]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
- ϕw = @view ϕ[indw]
-
- u_ux = Au_ux * uₕ + yAu_ux # u at ux
- ū_ux = Iu_ux * ϕu + yIu_ux # ū at ux
- ∂uū∂x = Cux * (u_ux .* ū_ux)
-
- u_uy = Au_uy * uₕ + yAu_uy # u at uy
- v̄_uy = Iv_uy * ϕv + yIv_uy # v̄ at uy
- ∂uv̄∂y = Cuy * (u_uy .* v̄_uy)
-
- u_uz = Au_uz * uₕ + yAu_uz # u at uz
- w̄_uz = Iw_uz * ϕw + yIw_uz # ū at uz
- ∂uw̄∂z = Cuz * (u_uz .* w̄_uz)
-
- v_vx = Av_vx * vₕ + yAv_vx # v at vx
- ū_vx = Iu_vx * ϕu + yIu_vx # ū at vx
- ∂vū∂x = Cvx * (v_vx .* ū_vx)
-
- v_vy = Av_vy * vₕ + yAv_vy # v at vy
- v̄_vy = Iv_vy * ϕv + yIv_vy # v̄ at vy
- ∂vv̄∂y = Cvy * (v_vy .* v̄_vy)
-
- v_vz = Av_vz * vₕ + yAv_vz # v at vz
- w̄_vz = Iw_vz * ϕw + yIw_vz # w̄ at vz
- ∂vw̄∂z = Cvz * (v_vz .* w̄_vz)
-
- w_wx = Aw_wx * wₕ + yAw_wx # w at wx
- ū_wx = Iu_wx * ϕu + yIu_wx # ū at wx
- ∂wū∂x = Cwx * (w_wx .* ū_wx)
-
- w_wy = Aw_wy * wₕ + yAw_wy # w at wy
- v̄_wy = Iv_wy * ϕv + yIv_wy # v̄ at wy
- ∂wv̄∂y = Cwy * (w_wy .* v̄_wy)
-
- w_wz = Aw_wz * wₕ + yAw_wz # w at wz
- w̄_wz = Iw_wz * ϕw + yIw_wz # w̄ at wz
- ∂ww̄∂z = Cwz * (w_wz .* w̄_wz)
-
- cu = @. ∂uū∂x + ∂uv̄∂y + ∂uw̄∂z
- cv = @. ∂vū∂x + ∂vv̄∂y + ∂vw̄∂z
- cw = @. ∂wū∂x + ∂wv̄∂y + ∂ww̄∂z
-
- c = [cu; cv; cw]
-
- if get_jacobian
- ## Convective terms, u-component
- C1 = Cux * Diagonal(ū_ux)
- C2 = Cux * Diagonal(u_ux) * newton_factor
- Conv_ux_11 = C1 * Au_ux .+ C2 * Iu_ux
-
- C1 = Cuy * Diagonal(v̄_uy)
- C2 = Cuy * Diagonal(u_uy) * newton_factor
- Conv_uy_11 = C1 * Au_uy
- Conv_uy_12 = C2 * Iv_uy
-
- C1 = Cuz * Diagonal(w̄_uz)
- C2 = Cuz * Diagonal(u_uz) * newton_factor
- Conv_uz_11 = C1 * Au_uz
- Conv_uz_13 = C2 * Iw_uz
-
- ## Convective terms, v-component
- C1 = Cvx * Diagonal(ū_vx)
- C2 = Cvx * Diagonal(v_vx) * newton_factor
- Conv_vx_21 = C2 * Iu_vx
- Conv_vx_22 = C1 * Av_vx
-
- C1 = Cvy * Diagonal(v̄_vy)
- C2 = Cvy * Diagonal(v_vy) * newton_factor
- Conv_vy_22 = C1 * Av_vy .+ C2 * Iv_vy
-
- C1 = Cvz * Diagonal(w̄_vz)
- C2 = Cvz * Diagonal(v_vz) * newton_factor
- Conv_vz_23 = C2 * Iw_vz
- Conv_vz_22 = C1 * Av_vz
-
- ## Convective terms, w-component
- C1 = Cwx * Diagonal(ū_wx)
- C2 = Cwx * Diagonal(w_wx) * newton_factor
- Conv_wx_31 = C2 * Iu_wx
- Conv_wx_33 = C1 * Aw_wx
-
- C1 = Cwy * Diagonal(v̄_wy)
- C2 = Cwy * Diagonal(w_wy) * newton_factor
- Conv_wy_32 = C2 * Iv_wy
- Conv_wy_33 = C1 * Aw_wy
-
- C1 = Cwz * Diagonal(w̄_wz)
- C2 = Cwz * Diagonal(w_wz) * newton_factor
- Conv_wz_33 = C1 * Aw_wz .+ C2 * Iw_wz
-
- ## Jacobian
- ∇c = [
- (Conv_ux_11+Conv_uy_11+Conv_uz_11) Conv_uy_12 Conv_uz_13
- Conv_vx_21 (Conv_vx_22+Conv_vy_22+Conv_vz_22) Conv_vz_23
- Conv_wx_31 Conv_wy_32 (Conv_wx_33+Conv_wy_33+Conv_wz_33)
- ]
- else
- ∇c = nothing
- end
-
- c, ∇c
-end
-
-"""
- convection_components!(
- c, ∇c, V, ϕ, setup, cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
- order4 = false,
- )
-
-Compute convection components.
-
-Mutating/non-allocating/in-place version.
-
-See also [`convection_components`](@ref).
-"""
-function convection_components! end
-
-convection_components!(c, ∇c, V, ϕ, setup, cache; kwargs...) =
- convection_components!(setup.grid.dimension, c, ∇c, V, ϕ, setup, cache; kwargs...)
-
-# 2D version
-function convection_components!(
- ::Dimension{2},
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
- order4 = false,
-)
- (; grid, operators) = setup
-
- if order4
- Cux = operators.Cux3
- Cuy = operators.Cuy3
- Cvx = operators.Cvx3
- Cvy = operators.Cvy3
-
- Au_ux = operators.Au_ux3
- Au_uy = operators.Au_uy3
- Av_vx = operators.Av_vx3
- Av_vy = operators.Av_vy3
-
- Iu_ux = operators.Iu_ux3
- Iv_uy = operators.Iv_uy3
- Iu_vx = operators.Iu_vx3
- Iv_vy = operators.Iv_vy3
-
- yAu_ux = bc_vectors.yAu_ux3
- yAu_uy = bc_vectors.yAu_uy3
- yAv_vx = bc_vectors.yAv_vx3
- yAv_vy = bc_vectors.yAv_vy3
-
- yIu_ux = bc_vectors.yIu_ux3
- yIv_uy = bc_vectors.yIv_uy3
- yIu_vx = bc_vectors.yIu_vx3
- yIv_vy = bc_vectors.yIv_vy3
- else
- (; Cux, Cuy, Cvx, Cvy) = operators
- (; Au_ux, Au_uy, Av_vx, Av_vy) = operators
- (; Iu_ux, Iv_uy, Iu_vx, Iv_vy) = operators
- (; yAu_ux, yAu_uy, yAv_vx, yAv_vy) = bc_vectors
- (; yIu_ux, yIv_uy, yIu_vx, yIv_vy) = bc_vectors
- end
-
- (; indu, indv) = grid
-
- (; u_ux, ū_ux, uū_ux, u_uy, v̄_uy, uv̄_uy) = cache
- (; v_vx, ū_vx, vū_vx, v_vy, v̄_vy, vv̄_vy) = cache
-
- (; ∂uū∂x, ∂uv̄∂y, ∂vū∂x, ∂vv̄∂y) = cache
- (; Conv_ux_11, Conv_uy_11, Conv_uy_12, Conv_vx_21, Conv_vx_22, Conv_vy_22) = cache
-
- cu = @view c[indu]
- cv = @view c[indv]
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
-
- # Convection components
- if order4
- # TODO: preallocated arrays for order4
- u_ux = Au_ux * uₕ + yAu_ux # U at ux
- ū_ux = Iu_ux * ϕu + yIu_ux # Ū at ux
- # ∂uū∂x = Cux * (u_ux .* ū_ux)
- uū_ux = @. u_ux = u_ux * ū_ux
-
- u_uy = Au_uy * uₕ + yAu_uy # U at uy
- v̄_uy = Iv_uy * ϕv + yIv_uy # Ū at uy
- # ∂uv̄∂y = Cuy * (u_uy .* v̄_uy)
- uv̄_uy = @. u_uy = u_uy * v̄_uy
-
- v_vx = Av_vx * vₕ + yAv_vx # V at vx
- ū_vx = Iu_vx * ϕu + yIu_vx # Ū at vx
- ∂vū∂x = Cvx * (v_vx .* ū_vx)
- vū_vx = @. v_vx = v_vx * ū_vx
-
- v_vy = Av_vy * vₕ + yAv_vy # V at vy
- v̄_vy = Iv_vy * ϕv + yIv_vy # Ū at vy
- # ∂vv̄∂y = Cvy * (v_vy .* v̄_vy)
- vv̄_vy = @. v_vy = v_vy * v̄_vy
- else
- # Fill preallocated arrays
- mul!(u_ux, Au_ux, uₕ)
- mul!(ū_ux, Iu_ux, ϕu)
- mul!(u_uy, Au_uy, uₕ)
- mul!(v̄_uy, Iv_uy, ϕv)
- mul!(v_vx, Av_vx, vₕ)
- mul!(ū_vx, Iu_vx, ϕu)
- mul!(v_vy, Av_vy, vₕ)
- mul!(v̄_vy, Iv_vy, ϕv)
-
- u_ux .+= yAu_ux
- ū_ux .+= yIu_ux
- @. uū_ux = u_ux * ū_ux
-
- u_uy .+= yAu_uy
- v̄_uy .+= yIv_uy
- @. uv̄_uy = u_uy * v̄_uy
-
- v_vx .+= yAv_vx
- ū_vx .+= yIu_vx
- @. vū_vx = v_vx * ū_vx
-
- v_vy .+= yAv_vy
- v̄_vy .+= yIv_vy
- @. vv̄_vy = v_vy * v̄_vy
- end
-
- mul!(∂uū∂x, Cux, uū_ux)
- mul!(∂uv̄∂y, Cuy, uv̄_uy)
- mul!(∂vū∂x, Cvx, vū_vx)
- mul!(∂vv̄∂y, Cvy, vv̄_vy)
-
- @. cu = ∂uū∂x + ∂uv̄∂y
- @. cv = ∂vū∂x + ∂vv̄∂y
-
- if get_jacobian
- ## Convective terms, u-component
- C1 = Cux * Diagonal(ū_ux)
- C2 = Cux * Diagonal(u_ux) * newton_factor
- Conv_ux_11 .= C1 * Au_ux .+ C2 * Iu_ux
-
- C1 = Cuy * Diagonal(v̄_uy)
- C2 = Cuy * Diagonal(u_uy) * newton_factor
- Conv_uy_11 .= C1 * Au_uy
- Conv_uy_12 .= C2 * Iv_uy
-
- ## Convective terms, v-component
- C1 = Cvx * Diagonal(ū_vx)
- C2 = Cvx * Diagonal(v_vx) * newton_factor
- Conv_vx_21 .= C2 * Iu_vx
- Conv_vx_22 .= C1 * Av_vx
-
- C1 = Cvy * Diagonal(v̄_vy)
- C2 = Cvy * Diagonal(v_vy) * newton_factor
- Conv_vy_22 .= C1 * Av_vy .+ C2 * Iv_vy
-
- ∇c[indu, indu] = Conv_ux_11 + Conv_uy_11
- ∇c[indu, indv] = Conv_uy_12
- ∇c[indv, indu] = Conv_vx_21
- ∇c[indv, indv] = Conv_vx_22 + Conv_vy_22
- end
-
- c, ∇c
-end
-
-# 3D version
-function convection_components!(
- ::Dimension{3},
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian = false,
- newton_factor = false,
- order4 = false,
-)
- order4 && error("order4 not implemented for 3D")
-
- (; grid, operators) = setup
- (; indu, indv, indw) = grid
- (; Cux, Cuy, Cuz, Cvx, Cvy, Cvz, Cwx, Cwy, Cwz) = operators
-
- (; Au_ux, Au_uy, Au_uz) = operators
- (; Av_vx, Av_vy, Av_vz) = operators
- (; Aw_wx, Aw_wy, Aw_wz) = operators
-
- (; Iu_ux, Iv_uy, Iw_uz) = operators
- (; Iu_vx, Iv_vy, Iw_vz) = operators
- (; Iu_wx, Iv_wy, Iw_wz) = operators
-
- (; yAu_ux, yAu_uy, yAu_uz) = bc_vectors
- (; yAv_vx, yAv_vy, yAv_vz) = bc_vectors
- (; yAw_wx, yAw_wy, yAw_wz) = bc_vectors
-
- (; yIu_ux, yIv_uy, yIw_uz) = bc_vectors
- (; yIu_vx, yIv_vy, yIw_vz) = bc_vectors
- (; yIu_wx, yIv_wy, yIw_wz) = bc_vectors
-
- (; u_ux, ū_ux, uū_ux, u_uy, v̄_uy, uv̄_uy, u_uz, w̄_uz, uw̄_uz) = cache
- (; v_vx, ū_vx, vū_vx, v_vy, v̄_vy, vv̄_vy, v_vz, w̄_vz, vw̄_vz) = cache
- (; w_wx, ū_wx, wū_wx, w_wy, v̄_wy, wv̄_wy, w_wz, w̄_wz, ww̄_wz) = cache
- (; ∂uū∂x, ∂uv̄∂y, ∂uw̄∂z, ∂vū∂x, ∂vv̄∂y, ∂vw̄∂z, ∂wū∂x, ∂wv̄∂y, ∂ww̄∂z) = cache
- (; Conv_ux_11, Conv_uy_11, Conv_uz_11, Conv_uy_12, Conv_uz_13) = cache
- (; Conv_vx_21, Conv_vx_22, Conv_vy_22, Conv_vz_22, Conv_vz_23) = cache
- (; Conv_wx_31, Conv_wy_32, Conv_wx_33, Conv_wy_33, Conv_wz_33) = cache
-
- cu = @view c[indu]
- cv = @view c[indv]
- cw = @view c[indw]
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- wₕ = @view V[indw]
-
- ϕu = @view ϕ[indu]
- ϕv = @view ϕ[indv]
- ϕw = @view ϕ[indw]
-
- # Convection components
- mul!(u_ux, Au_ux, uₕ)
- mul!(ū_ux, Iu_ux, ϕu)
- mul!(u_uy, Au_uy, uₕ)
- mul!(v̄_uy, Iv_uy, ϕv)
- mul!(u_uz, Au_uz, uₕ)
- mul!(w̄_uz, Iw_uz, ϕw)
-
- mul!(v_vx, Av_vx, vₕ)
- mul!(ū_vx, Iu_vx, ϕu)
- mul!(v_vy, Av_vy, vₕ)
- mul!(v̄_vy, Iv_vy, ϕv)
- mul!(v_vz, Av_vz, vₕ)
- mul!(w̄_vz, Iw_vz, ϕw)
-
- mul!(w_wx, Aw_wx, wₕ)
- mul!(ū_wx, Iu_wx, ϕu)
- mul!(w_wy, Aw_wy, wₕ)
- mul!(v̄_wy, Iv_wy, ϕv)
- mul!(w_wz, Aw_wz, wₕ)
- mul!(w̄_wz, Iw_wz, ϕw)
-
- u_ux .+= yAu_ux
- ū_ux .+= yIu_ux
- @. uū_ux = u_ux * ū_ux
-
- u_uy .+= yAu_uy
- v̄_uy .+= yIv_uy
- @. uv̄_uy = u_uy * v̄_uy
-
- u_uz .+= yAu_uz
- w̄_uz .+= yIw_uz
- @. uw̄_uz = u_uz * w̄_uz
-
- v_vx .+= yAv_vx
- ū_vx .+= yIu_vx
- @. vū_vx = v_vx * ū_vx
-
- v_vy .+= yAv_vy
- v̄_vy .+= yIv_vy
- @. vv̄_vy = v_vy * v̄_vy
-
- v_vz .+= yAv_vz
- w̄_vz .+= yIw_vz
- @. vw̄_vz = v_vz * w̄_vz
-
- w_wx .+= yAw_wx
- ū_wx .+= yIu_wx
- @. wū_wx = w_wx * ū_wx
-
- w_wy .+= yAw_wy
- v̄_wy .+= yIv_wy
- @. wv̄_wy = w_wy * v̄_wy
-
- w_wz .+= yAw_wz
- w̄_wz .+= yIw_wz
- @. ww̄_wz = w_wz * w̄_wz
-
- mul!(∂uū∂x, Cux, uū_ux)
- mul!(∂uv̄∂y, Cuy, uv̄_uy)
- mul!(∂uw̄∂z, Cuz, uw̄_uz)
-
- mul!(∂vū∂x, Cvx, vū_vx)
- mul!(∂vv̄∂y, Cvy, vv̄_vy)
- mul!(∂vw̄∂z, Cvz, vw̄_vz)
-
- mul!(∂wū∂x, Cwx, wū_wx)
- mul!(∂wv̄∂y, Cwy, wv̄_wy)
- mul!(∂ww̄∂z, Cwz, ww̄_wz)
-
- @. cu = ∂uū∂x + ∂uv̄∂y + ∂uw̄∂z
- @. cv = ∂vū∂x + ∂vv̄∂y + ∂vw̄∂z
- @. cw = ∂wū∂x + ∂wv̄∂y + ∂ww̄∂z
-
- if get_jacobian
- ## Convective terms, u-component
- C1 = Cux * Diagonal(ū_ux)
- C2 = Cux * Diagonal(u_ux) * newton_factor
- Conv_ux_11 .= C1 * Au_ux .+ C2 * Iu_ux
-
- C1 = Cuy * Diagonal(v̄_uy)
- C2 = Cuy * Diagonal(u_uy) * newton_factor
- Conv_uy_11 .= C1 * Au_uy
- Conv_uy_12 .= C2 * Iv_uy
-
- C1 = Cuz * Diagonal(w̄_uz)
- C2 = Cuz * Diagonal(u_uz) * newton_factor
- Conv_uz_11 .= C1 * Au_uz
- Conv_uz_13 .= C2 * Iw_uz
-
- ## Convective terms, v-component
- C1 = Cvx * Diagonal(ū_vx)
- C2 = Cvx * Diagonal(v_vx) * newton_factor
- Conv_vx_21 .= C2 * Iu_vx
- Conv_vx_22 .= C1 * Av_vx
-
- C1 = Cvy * Diagonal(v̄_vy)
- C2 = Cvy * Diagonal(v_vy) * newton_factor
- Conv_vy_22 .= C1 * Av_vy .+ C2 * Iv_vy
-
- C1 = Cvz * Diagonal(w̄_vz)
- C2 = Cvz * Diagonal(v_vz) * newton_factor
- Conv_vz_23 .= C2 * Iw_vz
- Conv_vz_22 .= C1 * Av_vz
-
- ## Convective terms, w-component
- C1 = Cwx * Diagonal(ū_wx)
- C2 = Cwx * Diagonal(w_wx) * newton_factor
- Conv_wx_31 .= C2 * Iu_wx
- Conv_wx_33 .= C1 * Aw_wx
-
- C1 = Cwy * Diagonal(v̄_wy)
- C2 = Cwy * Diagonal(w_wy) * newton_factor
- Conv_wy_32 .= C2 * Iv_wy
- Conv_wy_33 .= C1 * Aw_wy
-
- C1 = Cwz * Diagonal(w̄_wz)
- C2 = Cwz * Diagonal(w_wz) * newton_factor
- Conv_wz_33 .= C1 * Aw_wz .+ C2 * Iw_wz
-
- ## Jacobian
- ∇c[indu, indu] = Conv_ux_11 + Conv_uy_11 + Conv_uz_11
- ∇c[indu, indv] = Conv_uy_12
- ∇c[indu, indw] = Conv_uz_13
- ∇c[indv, indu] = Conv_vx_21
- ∇c[indv, indv] = Conv_vx_22 + Conv_vy_22 + Conv_vz_22
- ∇c[indv, indw] = Conv_vz_23
- ∇c[indw, indu] = Conv_wx_31
- ∇c[indw, indv] = Conv_wy_32
- ∇c[indw, indw] = Conv_wx_33 + Conv_wy_33 + Conv_wz_33
- end
-
- c, ∇c
-end
diff --git a/src/momentum/diffusion.jl b/src/momentum/diffusion.jl
deleted file mode 100644
index 3b7a66421..000000000
--- a/src/momentum/diffusion.jl
+++ /dev/null
@@ -1,233 +0,0 @@
-"""
- diffusion!(model, V, setup; bc_vectors, get_jacobian = false)
-
-Evaluate diffusive terms `d` and optionally Jacobian `∇d = ∂d/∂V` using viscosity model `model`.
-
-Non-mutating/allocating/out-of-place version.
-
-See also [`diffusion!`](@ref).
-"""
-function diffusion end
-
-diffusion(m, V, setup; kwargs...) = diffusion(setup.grid.dimension, m, V, setup; kwargs...)
-
-function diffusion(m::LaminarModel, V, setup; bc_vectors, get_jacobian = false)
- (; Re) = m
- (; Diff) = setup.operators
- (; yDiff) = bc_vectors
-
- d = 1 ./ Re .* (Diff * V .+ yDiff)
-
- if get_jacobian
- ∇d = 1 / Re * Diff
- else
- ∇d = nothing
- end
-
- d, ∇d
-end
-
-# 2D version
-function diffusion(
- ::Dimension{2},
- model::Union{QRModel,SmagorinskyModel,MixingLengthModel},
- V,
- setup;
- bc_vectors,
- get_jacobian = false,
-)
- (; indu, indv) = setup.grid
- (; Dux, Duy, Dvx, Dvy) = setup.operators
- (; Su_ux, Su_uy, Su_vx, Sv_vx, Sv_vy, Sv_uy) = setup.operators
- (; Aν_ux, Aν_uy, Aν_vx, Aν_vy) = setup.operators
- (; yAν_ux, yAν_uy, yAν_vx, yAν_vy) = bc_vectors
-
- # Get components of strain tensor and its magnitude;
- # The magnitude S_abs is evaluated at pressure points
- S11, S12, S21, S22, S_abs, S_abs_u, S_abs_v =
- strain_tensor(V, setup; bc_vectors, get_jacobian)
-
- # Turbulent viscosity at all pressure points
- ν_t = turbulent_viscosity(model, setup, S_abs)
-
- # To compute the diffusion, we need ν_t at ux, uy, vx and vy locations
- # This means we have to reverse the process of `strain_tensor`: go
- # from pressure points back to the ux, uy, vx, vy locations
- ν_t_ux = Aν_ux * ν_t + yAν_ux
- ν_t_uy = Aν_uy * ν_t + yAν_uy
- ν_t_vx = Aν_vx * ν_t + yAν_vx
- ν_t_vy = Aν_vy * ν_t + yAν_vy
-
- # Now the total diffusive terms (laminar + turbulent) is as follows
- # Note that the factor 2 is because
- # Tau = 2*(ν+ν_t)*S(u), with S(u) = 1/2*(∇u + (∇u)^T)
-
- # Molecular viscosity
- ν = 1 / model.Re
-
- du = Dux * (2 .* (ν .+ ν_t_ux) .* S11[:]) .+ Duy * (2 .* (ν .+ ν_t_uy) .* S12[:])
- dv = Dvx * (2 .* (ν .+ ν_t_vx) .* S21[:]) .+ Dvy * (2 .* (ν .+ ν_t_vy) .* S22[:])
-
- if get_jacobian
- # Freeze ν_t, i.e. we skip the derivative of ν_t wrt V in the Jacobian
- Jacu1 =
- Dux * 2 * spdiagm(ν .+ ν_t_ux) * Su_ux +
- Duy * 2 * spdiagm(ν .+ ν_t_uy) * 1 / 2 * Su_uy
- Jacu2 = Duy * 2 * spdiagm(ν .+ ν_t_uy) * 1 / 2 * Sv_uy
- Jacv1 = Dvx * 2 * spdiagm(ν .+ ν_t_vx) * 1 / 2 * Su_vx
- Jacv2 =
- Dvx * 2 * spdiagm(ν .+ ν_t_vx) * 1 / 2 * Sv_vx +
- Dvy * 2 * spdiagm(ν .+ ν_t_vy) * Sv_vy
- Jacu = [Jacu1 Jacu2]
- Jacv = [Jacv1 Jacv2]
-
- K = turbulent_K(model, setup)
-
- tmpu1 =
- 2 * Dux * spdiagm(S11) * Aν_ux * S_abs_u +
- 2 * Duy * spdiagm(S12) * Aν_uy * S_abs_u
- tmpu2 = 2 * Duy * spdiagm(S12) * Aν_uy * S_abs_v
- tmpv1 = 2 * Dvx * spdiagm(S21) * Aν_vx * S_abs_u
- tmpv2 =
- 2 * Dvx * spdiagm(S21) * Aν_vx * S_abs_v +
- 2 * Dvy * spdiagm(S22) * Aν_vy * S_abs_v
- Jacu += K * [tmpu1 tmpu2]
- Jacv += K * [tmpv1 tmpv2]
-
- ∇d = [Jacu; Jacv]
- else
- ∇d = nothing
- end
-
- d = [du; dv]
-
- d, ∇d
-end
-
-# 3D version
-function diffusion(
- ::Dimension{3},
- model::Union{QRModel,SmagorinskyModel,MixingLengthModel},
- V,
- setup;
- bc_vectors,
- get_jacobian = false,
-)
- error("Not implemented")
-end
-
-"""
- diffusion!(model, d, ∇d, V, setup; bc_vectors, get_jacobian = false)
-
-Evaluate diffusive terms `d` and optionally Jacobian `∇d = ∂d/∂V` using viscosity model `model`.
-"""
-function diffusion! end
-
-diffusion!(m, d, ∇d, V, setup; kwargs...) =
- diffusion!(setup.grid.dimension, m, d, ∇d, V, setup; kwargs...)
-
-function diffusion!(m::LaminarModel, d, ∇d, V, setup; bc_vectors, get_jacobian = false)
- (; Re) = m
- (; Diff) = setup.operators
- (; yDiff) = bc_vectors
-
- # d = 1 ./ Re .* (Diff * V .+ yDiff)
- mul!(d, Diff, V)
- @. d = 1 / Re * (d + yDiff)
-
- get_jacobian && (@. ∇d = 1 / Re * Diff)
-
- d, ∇d
-end
-
-# 2D version
-function diffusion!(
- ::Dimension{2},
- model::Union{QRModel,SmagorinskyModel,MixingLengthModel},
- d,
- ∇d,
- V,
- setup;
- bc_vectors,
- get_jacobian = false,
-)
- (; indu, indv) = setup.grid
- (; Dux, Duy, Dvx, Dvy) = setup.operators
- (; Su_ux, Su_uy, Su_vx, Sv_vx, Sv_vy, Sv_uy) = setup.operators
- (; Aν_ux, Aν_uy, Aν_vx, Aν_vy) = setup.operators
- (; yAν_ux, yAν_uy, yAν_vx, yAν_vy) = bc_vectors
-
- du = @view d[indu]
- dv = @view d[indv]
-
- # Get components of strain tensor and its magnitude;
- # The magnitude S_abs is evaluated at pressure points
- S11, S12, S21, S22, S_abs, S_abs_u, S_abs_v =
- strain_tensor(V, setup; bc_vectors, get_jacobian)
-
- # Turbulent viscosity at all pressure points
- ν_t = turbulent_viscosity(model, setup, S_abs)
-
- # To compute the diffusion, we need ν_t at ux, uy, vx and vy locations
- # This means we have to reverse the process of `strain_tensor`: go
- # from pressure points back to the ux, uy, vx, vy locations
- ν_t_ux = Aν_ux * ν_t + yAν_ux
- ν_t_uy = Aν_uy * ν_t + yAν_uy
- ν_t_vx = Aν_vx * ν_t + yAν_vx
- ν_t_vy = Aν_vy * ν_t + yAν_vy
-
- # Now the total diffusive terms (laminar + turbulent) is as follows
- # Note that the factor 2 is because
- # Tau = 2*(ν+ν_t)*S(u), with S(u) = 1/2*(∇u + (∇u)^T)
-
- # Molecular viscosity
- ν = 1 / model.Re
-
- du .= Dux * (2 .* (ν .+ ν_t_ux) .* S11) .+ Duy * (2 .* (ν .+ ν_t_uy) .* S12)
- dv .= Dvx * (2 .* (ν .+ ν_t_vx) .* S21) .+ Dvy * (2 .* (ν .+ ν_t_vy) .* S22)
-
- if get_jacobian
- # Freeze ν_t, i.e. we skip the derivative of ν_t wrt V in the Jacobian
- Jacu1 =
- Dux * 2 * spdiagm(ν .+ ν_t_ux) * Su_ux +
- Duy * 2 * spdiagm(ν .+ ν_t_uy) * 1 // 2 * Su_uy
- Jacu2 = Duy * 2 * spdiagm(ν .+ ν_t_uy) * 1 // 2 * Sv_uy
- Jacv1 = Dvx * 2 * spdiagm(ν .+ ν_t_vx) * 1 // 2 * Su_vx
- Jacv2 =
- Dvx * 2 * spdiagm(ν .+ ν_t_vx) * 1 // 2 * Sv_vx +
- Dvy * 2 * spdiagm(ν .+ ν_t_vy) * Sv_vy
- Jacu = [Jacu1 Jacu2]
- Jacv = [Jacv1 Jacv2]
-
- K = turbulent_K(model, setup)
-
- tmpu1 =
- 2 * Dux * spdiagm(S11) * Aν_ux * S_abs_u +
- 2 * Duy * spdiagm(S12) * Aν_uy * S_abs_u
- tmpu2 = 2 * Duy * spdiagm(S12) * Aν_uy * S_abs_v
- tmpv1 = 2 * Dvx * spdiagm(S21) * Aν_vx * S_abs_u
- tmpv2 =
- 2 * Dvx * spdiagm(S21) * Aν_vx * S_abs_v +
- 2 * Dvy * spdiagm(S22) * Aν_vy * S_abs_v
- Jacu += K * [tmpu1 tmpu2]
- Jacv += K * [tmpv1 tmpv2]
-
- ∇d .= [Jacu; Jacv]
- end
-
- d, ∇d
-end
-
-# 3D version
-function diffusion!(
- ::Dimension{3},
- model::Union{QRModel,SmagorinskyModel,MixingLengthModel},
- d,
- ∇d,
- V,
- setup;
- bc_vectors,
- get_jacobian = false,
-)
- error("Not implemented")
-end
diff --git a/src/momentum/momentum.jl b/src/momentum/momentum.jl
deleted file mode 100644
index eef27c268..000000000
--- a/src/momentum/momentum.jl
+++ /dev/null
@@ -1,155 +0,0 @@
-"""
- momentum(
- V, ϕ, p, t, setup;
- bc_vectors = nothing,
- get_jacobian = false,
- nopressure = false,
- newton_factor = false,
- )
-
-Calculate RHS of momentum equations and, optionally, Jacobian with respect to velocity field.
-
- - `V`: velocity field
- - `ϕ`: convected field: e.g. ``\\frac{\\partial (\\phi_x V)}{\\partial x} + \\frac{\\partial (\\phi_y V)}{\\partial y}``; usually `ϕ = V` (so `ϕx = u`, `ϕy = v`)
- - `p`: pressure
- - `bc_vectors`: boundary condition vectors `y`
- - `get_jacobian`: return `∇F = ∂F/∂V`
- - `nopressure`: exclude pressure gradient; in this case input argument `p` is not used
- - `newton_factor`
-
-Non-mutating/allocating/out-of-place version.
-
-See also [`momentum!`](@ref).
-"""
-function momentum(
- V,
- ϕ,
- p,
- t,
- setup;
- bc_vectors = nothing,
- get_jacobian = false,
- nopressure = false,
- newton_factor = false,
-)
- (; viscosity_model, convection_model, force, boundary_conditions, operators) = setup
- (; G) = operators
-
- # Unsteady BC
- if isnothing(bc_vectors) || boundary_conditions.bc_unsteady
- bc_vectors = get_bc_vectors(setup, t)
- end
- (; y_p) = bc_vectors
-
- # Convection
- c, ∇c =
- convection(convection_model, V, ϕ, setup; bc_vectors, get_jacobian, newton_factor)
-
- # Diffusion
- d, ∇d = diffusion(viscosity_model, V, setup; bc_vectors, get_jacobian)
-
- # Body force
- b = force
-
- # Residual in Finite Volume form, including the pressure contribution
- F = @. -c + d + b
-
- # Nopressure = false is the most common situation, in which we return the entire
- # right-hand side vector
- if !nopressure
- F = F .- (G * p .+ y_p)
- end
-
- if get_jacobian
- # Jacobian requested
- # We return only the Jacobian with respect to V (not p)
- ∇F = @. -∇c + ∇d
- else
- ∇F = nothing
- end
-
- F, ∇F
-end
-
-"""
- momentum!(F, ∇F, V, ϕ, p, t, setup, cache; get_jacobian = false, nopressure = false)
-
-Calculate rhs of momentum equations and, optionally, Jacobian with respect to velocity field.
-
- - `V`: velocity field
- - `ϕ`: convected field: e.g. ``\\frac{\\partial (\\phi_x V)}{\\partial x} + \\frac{\\partial (\\phi_y V)}{\\partial y}``; usually `ϕ = V` (so `ϕx = u`, `ϕy = v`)
- - `p`: pressure
- - `bc_vectors`: boundary condition vectors `y`
- - `get_jacobian`: return `∇F = ∂F/∂V`
- - `nopressure`: exclude pressure gradient; in this case input argument `p` is not used
- - `newton_factor`
-
-Mutating/non-allocating/in-place version.
-
-See also [`momentum`](@ref).
-"""
-function momentum!(
- F,
- ∇F,
- V,
- ϕ,
- p,
- t,
- setup,
- cache;
- bc_vectors = nothing,
- get_jacobian = false,
- nopressure = false,
- newton_factor = false,
-)
- (; viscosity_model, convection_model, force, boundary_conditions, operators) = setup
- (; G) = setup.operators
-
- # Unsteady BC
- if isnothing(bc_vectors) || boundary_conditions.bc_unsteady
- bc_vectors = get_bc_vectors(setup, t)
- end
- (; y_p) = bc_vectors
-
- # Store intermediate results in temporary variables
- (; c, ∇c, d, ∇d, b, Gp) = cache
-
- # Convection
- convection!(
- convection_model,
- c,
- ∇c,
- V,
- ϕ,
- setup,
- cache;
- bc_vectors,
- get_jacobian,
- newton_factor,
- )
-
- # Diffusion
- diffusion!(viscosity_model, d, ∇d, V, setup; bc_vectors, get_jacobian)
-
- # Body force
- b = force
-
- # Residual in Finite Volume form, including the pressure contribution
- @. F = -c + d + b
-
- # Nopressure = false is the most common situation, in which we return the entire
- # right-hand side vector
- if !nopressure
- mul!(Gp, G, p)
- Gp .+= y_p
- @. F -= Gp
- end
-
- if get_jacobian
- # Jacobian requested
- # We return only the Jacobian with respect to V (not p)
- @. ∇F = -∇c + ∇d
- end
-
- F, ∇F
-end
diff --git a/src/momentum/momentumcache.jl b/src/momentum/momentumcache.jl
deleted file mode 100644
index 441d51a07..000000000
--- a/src/momentum/momentumcache.jl
+++ /dev/null
@@ -1,242 +0,0 @@
-"""
- MomentumCache(setup)
-
-Preallocation structure for terms in the momentum equations.
-"""
-MomentumCache(setup, V, p) = MomentumCache(setup.grid.dimension, setup, V, p)
-
-function MomentumCache(::Dimension{2}, setup, V, p)
- (; grid, operators) = setup
- (; Nu, Nv, NV) = grid
- (; Iu_ux, Iv_uy, Iu_vx, Iv_vy) = operators
-
- T = eltype(V)
-
- Gp = similar(V)
- c = similar(V)
- c2 = similar(V)
- c3 = similar(V)
- d = similar(V)
- b = similar(V)
-
- ∇c = spzeros(T, NV, NV)
- ∇c2 = spzeros(T, NV, NV)
- ∇c3 = spzeros(T, NV, NV)
- ∇d = spzeros(T, NV, NV)
- ∇b = spzeros(T, NV, NV)
-
- u_ux = similar(V, size(Iu_ux, 1))
- ū_ux = similar(V, size(Iu_ux, 1))
- uū_ux = similar(V, size(Iu_ux, 1))
-
- u_uy = similar(V, size(Iv_uy, 1))
- v̄_uy = similar(V, size(Iv_uy, 1))
- uv̄_uy = similar(V, size(Iv_uy, 1))
-
- v_vx = similar(V, size(Iu_vx, 1))
- ū_vx = similar(V, size(Iu_vx, 1))
- vū_vx = similar(V, size(Iu_vx, 1))
-
- v_vy = similar(V, size(Iv_vy, 1))
- v̄_vy = similar(V, size(Iv_vy, 1))
- vv̄_vy = similar(V, size(Iv_vy, 1))
-
- ∂uū∂x = similar(V, Nu)
- ∂uv̄∂y = similar(V, Nu)
- ∂vū∂x = similar(V, Nv)
- ∂vv̄∂y = similar(V, Nv)
-
- Conv_ux_11 = spzeros(T, Nu, Nu)
- Conv_uy_11 = spzeros(T, Nu, Nu)
- Conv_uy_12 = spzeros(T, Nu, Nv)
-
- Conv_vx_21 = spzeros(T, Nv, Nu)
- Conv_vx_22 = spzeros(T, Nv, Nv)
- Conv_vy_22 = spzeros(T, Nv, Nv)
-
- (;
- Gp,
- c,
- c2,
- c3,
- d,
- b,
- ∇c,
- ∇c2,
- ∇c3,
- ∇d,
- ∇b,
- u_ux,
- ū_ux,
- uū_ux,
- u_uy,
- v̄_uy,
- uv̄_uy,
- v_vx,
- ū_vx,
- vū_vx,
- v_vy,
- v̄_vy,
- vv̄_vy,
- ∂uū∂x,
- ∂uv̄∂y,
- ∂vū∂x,
- ∂vv̄∂y,
- Conv_ux_11,
- Conv_uy_11,
- Conv_uy_12,
- Conv_vx_21,
- Conv_vx_22,
- Conv_vy_22,
- )
-end
-
-function MomentumCache(::Dimension{3}, setup, V, p)
- (; grid, operators) = setup
- (; Nu, Nv, Nw, NV) = grid
- (; Iu_ux, Iv_uy, Iw_uz, Iu_vx, Iv_vy, Iw_vz, Iu_wx, Iv_wy, Iw_wz) = operators
-
- T = eltype(Iu_ux)
-
- Gp = similar(V, NV)
- c = similar(V, NV)
- c2 = similar(V, NV)
- c3 = similar(V, NV)
- d = similar(V, NV)
- b = similar(V, NV)
-
- ∇c = spzeros(T, NV, NV)
- ∇c2 = spzeros(T, NV, NV)
- ∇c3 = spzeros(T, NV, NV)
- ∇d = spzeros(T, NV, NV)
- ∇b = spzeros(T, NV, NV)
-
- u_ux = similar(V, size(Iu_ux, 1))
- ū_ux = similar(V, size(Iu_ux, 1))
- uū_ux = similar(V, size(Iu_ux, 1))
-
- u_uy = similar(V, size(Iv_uy, 1))
- v̄_uy = similar(V, size(Iv_uy, 1))
- uv̄_uy = similar(V, size(Iv_uy, 1))
-
- u_uz = similar(V, size(Iw_uz, 1))
- w̄_uz = similar(V, size(Iw_uz, 1))
- uw̄_uz = similar(V, size(Iw_uz, 1))
-
- v_vx = similar(V, size(Iu_vx, 1))
- ū_vx = similar(V, size(Iu_vx, 1))
- vū_vx = similar(V, size(Iu_vx, 1))
-
- v_vy = similar(V, size(Iv_vy, 1))
- v̄_vy = similar(V, size(Iv_vy, 1))
- vv̄_vy = similar(V, size(Iv_vy, 1))
-
- v_vz = similar(V, size(Iw_vz, 1))
- w̄_vz = similar(V, size(Iw_vz, 1))
- vw̄_vz = similar(V, size(Iw_vz, 1))
-
- w_wx = similar(V, size(Iu_wx, 1))
- ū_wx = similar(V, size(Iu_wx, 1))
- wū_wx = similar(V, size(Iu_wx, 1))
-
- w_wy = similar(V, size(Iv_wy, 1))
- v̄_wy = similar(V, size(Iv_wy, 1))
- wv̄_wy = similar(V, size(Iv_wy, 1))
-
- w_wz = similar(V, size(Iw_wz, 1))
- w̄_wz = similar(V, size(Iw_wz, 1))
- ww̄_wz = similar(V, size(Iw_wz, 1))
-
- ∂uū∂x = similar(V, Nu)
- ∂uv̄∂y = similar(V, Nu)
- ∂uw̄∂z = similar(V, Nu)
- ∂vū∂x = similar(V, Nv)
- ∂vv̄∂y = similar(V, Nv)
- ∂vw̄∂z = similar(V, Nv)
- ∂wū∂x = similar(V, Nw)
- ∂wv̄∂y = similar(V, Nw)
- ∂ww̄∂z = similar(V, Nw)
-
- Conv_ux_11 = spzeros(T, Nu, Nu)
- Conv_uy_11 = spzeros(T, Nu, Nu)
- Conv_uz_11 = spzeros(T, Nu, Nu)
- Conv_uy_12 = spzeros(T, Nu, Nv)
- Conv_uz_13 = spzeros(T, Nu, Nw)
-
- Conv_vx_21 = spzeros(T, Nv, Nu)
- Conv_vx_22 = spzeros(T, Nv, Nv)
- Conv_vy_22 = spzeros(T, Nv, Nv)
- Conv_vz_22 = spzeros(T, Nv, Nv)
- Conv_vz_23 = spzeros(T, Nv, Nw)
-
- Conv_wx_31 = spzeros(T, Nw, Nu)
- Conv_wy_32 = spzeros(T, Nw, Nv)
- Conv_wx_33 = spzeros(T, Nw, Nw)
- Conv_wy_33 = spzeros(T, Nw, Nw)
- Conv_wz_33 = spzeros(T, Nw, Nw)
-
- (;
- Gp,
- c,
- c2,
- c3,
- d,
- b,
- ∇c,
- ∇c2,
- ∇c3,
- ∇d,
- ∇b,
- u_ux,
- ū_ux,
- uū_ux,
- u_uy,
- v̄_uy,
- uv̄_uy,
- u_uz,
- w̄_uz,
- uw̄_uz,
- v_vx,
- ū_vx,
- vū_vx,
- v_vy,
- v̄_vy,
- vv̄_vy,
- v_vz,
- w̄_vz,
- vw̄_vz,
- w_wx,
- ū_wx,
- wū_wx,
- w_wy,
- v̄_wy,
- wv̄_wy,
- w_wz,
- w̄_wz,
- ww̄_wz,
- ∂uū∂x,
- ∂uv̄∂y,
- ∂uw̄∂z,
- ∂vū∂x,
- ∂vv̄∂y,
- ∂vw̄∂z,
- ∂wū∂x,
- ∂wv̄∂y,
- ∂ww̄∂z,
- Conv_ux_11,
- Conv_uy_11,
- Conv_uz_11,
- Conv_uy_12,
- Conv_uz_13,
- Conv_vx_21,
- Conv_vx_22,
- Conv_vy_22,
- Conv_vz_22,
- Conv_vz_23,
- Conv_wx_31,
- Conv_wy_32,
- Conv_wx_33,
- Conv_wy_33,
- Conv_wz_33,
- )
-end
diff --git a/src/momentum/strain_tensor.jl b/src/momentum/strain_tensor.jl
deleted file mode 100644
index ddf462182..000000000
--- a/src/momentum/strain_tensor.jl
+++ /dev/null
@@ -1,163 +0,0 @@
-"""
- strain_tensor(V, setup; bc_vectors, get_jacobian = false, get_S_abs = false)
-
-Evaluate rate of strain tensor `S(V)` and its magnitude.
-"""
-function strain_tensor end
-
-strain_tensor(V, setup; kwargs...) =
- strain_tensor(setup.grid.dimension, V, setup; kwargs...)
-
-# 2D version
-function strain_tensor(
- ::Dimension{2},
- V,
- setup;
- bc_vectors,
- get_jacobian = false,
- get_S_abs = false,
-)
- (; grid, operators, boundary_conditions) = setup
- (; Nx, Ny, Nu, Nv, Np, indu, indv) = grid
- (; Nux_in, Nuy_in, Nvx_in, Nvy_in) = grid
- (; x, y, xp, yp) = grid
- (; Su_ux, Su_uy, Su_vx, Sv_vx, Sv_vy, Sv_uy) = operators
- (; Cux_k, Cuy_k, Cvx_k, Cvy_k, Auy_k, Avx_k) = operators
- (; ySu_ux, ySu_uy, ySu_vx, ySv_vx, ySv_vy, ySv_uy) = bc_vectors
- (; yCux_k, yCuy_k, yCvx_k, yCvy_k, yAuy_k, yAvx_k) = bc_vectors
-
- T = eltype(V)
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- # These four components are approximated by
- S11 = Su_ux * uₕ + ySu_ux
- S12 = 1 // 2 .* (Su_uy * uₕ .+ ySu_uy .+ Sv_uy * vₕ .+ ySv_uy)
- S21 = 1 // 2 .* (Su_vx * uₕ .+ ySu_vx .+ Sv_vx * vₕ .+ ySv_vx)
- S22 = Sv_vy * vₕ + ySv_vy
-
- # Note: S11 and S22 at xp, yp locations (pressure locations)
- # S12, S21 at vorticity locations (corners of pressure cells, (x, y))
-
- # Option 1: get each S11, S12, S21, S22 at 4 locations (ux locations, at uy locations, at vx
- # Locations and at vy locations); this gives 16 S fields. determine S_abs at each of these
- # Locations, giving 4 S_abs fields, that can be used in computing
- # Dux*(S_abs_ux .* (Su_ux*uₕ+ySu_ux)) etc.
-
- # Option 2: interpolate S11, S12, S21, S22 to pressure locations
- # Determine S_abs at pressure location
- # Then interpolate to ux, uy, vx, vy locations
-
- # We will use option 2;
- # Within option 2, we can decide to interpolate S12, S21 etc (option 2a), or we can use
- # Directly the operators that map from velocity field to the S locations,
- # As used for example in ke_production (option 2b).
-
- if get_S_abs
- # Option 2b
- if boundary_conditions.u.x == (:periodic, :periodic)
- # "cut-off" the double points in case of periodic BC
- # For periodic boundary conditions S11(Npx+1, :) = S11(1, :)
- # So S11 has size (Npx+1)*Npy; the last row are "ghost" points equal to the
- # First points. we have S11 at positions ([xp[1] - 1/2*(hx[1]+hx[end]); xp], yp)
- S11_p = reshape(S11, Nux_in + 1, Nuy_in)
- S11_p = S11_p(2:(Nux_in+1), :) # B
-
- # S12 is defined on the corners: size Nux_in*(Nuy_in+1), positions (xin, y)
- # Get S12 and S21 at all corner points
- S12_temp = zeros(T, Nx + 1, Ny + 1)
- S12_temp[1:Nx, :] = reshape(S12, Nx, Ny + 1)
- S12_temp[Nx+1, :] = S12_temp[1, :]
- elseif boundary_conditions.u.x[1] == :dirichlet &&
- boundary_conditions.u.x[2] == :pressure
- S11_p = reshape(S11, Nux_in + 1, Nuy_in)
- S11_p = S11_p[1:Nux_in, :] # Cut off last point
-
- # S12 is defined on the corners: size Nux_in*(Nuy_in+1), positions (xin, y)
- # Get S12 and S21 at all corner points
- S12_temp = zeros(T, Nx + 1, Ny + 1)
- S12_temp[2:(Nx+1), :] = reshape(S12, Nx, Ny + 1)
- S12_temp[1, :] = S12_temp[2, :] # Copy from x[2] to x[1]; one could do this more accurately in principle by using the BC
- else
- error("BC not implemented in strain_tensor.jl")
- end
-
- if boundary_conditions.v.y[1] == :periodic &&
- boundary_conditions.v.y[2] == :periodic
- # Similarly, S22(:, Npy+1) = S22(:, 1). positions (xp, [yp;yp[1]])
- S22_p = reshape(S22, Nvx_in, Nvy_in + 1)
- S22_p = S22_p(:, 2:(Nvy_in+1)) # Why not 1:Nvy_in?
-
- # Similarly S21 is size (Nux_in+1)*Nuy_in, positions (x, yin)
- S21_temp = zeros(T, Nx + 1, Ny + 1)
- S21_temp[:, 1:Ny] = reshape(S21, Nx + 1, Ny)
- S21_temp[:, Ny+1] = S21_temp[:, 1]
- elseif boundary_conditions.v.y[1] == :pressure &&
- boundary_conditions.v.y[2] == :pressure
- S22_p = reshape(S22, Nvx_in, Nvy_in + 1)
- S22_p = S22_p(:, 2:Nvy_in)
-
- # This is nicely defined on all corners
- S21_temp = reshape(S21, Nx + 1, Ny + 1)
- else
- error("BC not implemented in strain_tensor.jl")
- end
-
- # Now interpolate S12 and S21 to pressure points
- # S11 and S22 have already been trimmed down to this grid
-
- # S21 and S12 should be equal!
- # FIXME: Find interpolation syntax
- error("Interpolation not implemented")
- # S12_p = interp2(y', x, S12_temp, yp', xp)
- # S21_p = interp2(y', x, S21_temp, yp', xp)
-
- ## Invariants
- q = @. 1 // 2 * (S11_p[:]^2 + S12_p[:]^2 + S21_p[:]^2 + S22_p[:]^2)
-
- # Absolute value of strain tensor
- # With S as defined above, i.e. 1/2*(grad u + grad u^T)
- # S_abs = sqrt(2*tr(S^2)) = sqrt(4*q)
- S_abs = sqrt(4q)
- else
- # Option 2a
- S11_p = Cux_k * uₕ + yCux_k
- S12_p =
- 1 // 2 .* (
- Cuy_k * (Auy_k * uₕ + yAuy_k) .+ yCuy_k .+ Cvx_k * (Avx_k * vₕ + yAvx_k) .+
- yCvx_k
- )
- S21_p = S12_p
- S22_p = Cvy_k * vₕ + yCvy_k
-
- S_abs = @. sqrt(2 * S11_p^2 + 2 * S22_p^2 + 2 * S12_p^2 + 2 * S21_p^2)
-
- # Jacobian of S_abs wrt u and v
- if get_jacobian
- ϵ = 100 * eps(T)
- Sabs_inv = spdiagm(1 ./ (2 .* S_abs .+ ϵ))
- Jacu =
- Sabs_inv * (4 * spdiagm(S11_p) * Cux_k + 4 * spdiagm(S12_p) * Cuy_k * Auy_k)
- Jacv =
- Sabs_inv * (4 * spdiagm(S12_p) * Cvx_k * Avx_k + 4 * spdiagm(S22_p) * Cvy_k)
- else
- Jacu = spzeros(T, Np, Nu)
- Jacv = spzeros(T, Np, Nv)
- end
- end
-
- S11, S12, S21, S22, S_abs, Jacu, Jacv
-end
-
-# 3D version
-function strain_tensor(
- ::Dimension{3},
- V,
- setup;
- bc_vectors,
- get_jacobian = false,
- get_S_abs = false,
-)
- error("Not implemented (3D)")
-end
diff --git a/src/momentum/turbulent_K.jl b/src/momentum/turbulent_K.jl
deleted file mode 100644
index 394ff7573..000000000
--- a/src/momentum/turbulent_K.jl
+++ /dev/null
@@ -1,17 +0,0 @@
-"""
- turbulent_K(model, setup)
-
-Compute the constant part of the turbulent viscosity.
-"""
-function turbulent_K end
-
-turbulent_K(model::SmagorinskyModel, setup) = model.C_s^2 * max_size(setup.grid)^2
-
-function turbulent_K(::QRModel, setup)
- (; α_reg) = setup.operators
- Δ = max_size(setup.grid)
- C_d = Δ^2 / 8
- C_d * 1 // 2 * (1 - α_reg / C_d)^2
-end
-
-turbulent_K(model::MixingLengthModel, setup) = model.lm^2
diff --git a/src/momentum/turbulent_viscosity.jl b/src/momentum/turbulent_viscosity.jl
deleted file mode 100644
index 81e980729..000000000
--- a/src/momentum/turbulent_viscosity.jl
+++ /dev/null
@@ -1,8 +0,0 @@
-"""
- turbulent_viscosity(model, setup, S_abs)
-
-Compute turbulent viscosity based on `S_abs`.
-"""
-function turbulent_viscosity end
-
-turbulent_viscosity(model, setup, S_abs) = turbulent_K(model, setup) * S_abs
diff --git a/src/operators.jl b/src/operators.jl
new file mode 100644
index 000000000..59be232ed
--- /dev/null
+++ b/src/operators.jl
@@ -0,0 +1,1426 @@
+# Note on implementation:
+# This file contains various differential operators.
+#
+# Each operator comes with
+#
+# - an modifying in-place version, e.g. `divergence!(div, u, setup)`,
+# - an allocating out-of-place version, e.g. `div = divergence(u, setup)`.
+#
+# The out-of-place versions can be used as building blocks in a
+# Zygote-differentiable program, thanks to the `rrule` methods
+# defined.
+#
+# The domain is divided into `N = (N[1], ..., N[D])` finite volumes.
+# These also include ghost volumes, possibly outside the domain.
+# For a Cartesian index `I`, volume center fields are naturally in the center,
+# but volume face fields are always to the _right_ of volume I.
+#
+# _All_ fields have the size `N`. These `N` components include
+#
+# - degrees of freedom
+# - boundary values, which are still used, but are filled in separately
+# - unused values, which are never used at all. These are still there so that
+# we can guarantee that `ω[I]`, `u[1][I]`, `u[2][I]`, and `p[I]` etc. are
+# at their canonical position in to the volume `I`. Otherwise we would
+# need an offset for each BC type and each combination. Asymptotically
+# speaking (for large `N`), the additional memory footprint of having these
+# around is negligible.
+#
+# The operators are implemented as kernels.
+# The kernels are called for each index in `ndrange`, typically set
+# to the degrees of freedom of the output quantity. Boundary values for the
+# output quantity are filled in separately, by calling `apply_bc_*` when needed.
+# It is assumed that the appropriate boundary values for the input fields are
+# already filled in.
+#
+# The adjoint kernels are written manually for now.
+# In the future, Enzyme.jl might be able to do this automatically.
+
+"""
+ e = Offset{D}()
+
+Cartesian index unit vector in `D = 2` or `D = 3` dimensions.
+Calling `e(α)` returns a Cartesian index with `1` in the dimension `α` and zeros
+elsewhere.
+
+See
+for writing kernel loops using Cartesian indices.
+"""
+struct Offset{D} end
+
+@inline (::Offset{D})(α) where {D} = CartesianIndex(ntuple(β -> β == α ? 1 : 0, D))
+
+"""
+ divergence!(div, u, setup)
+
+Compute divergence of velocity field (in-place version).
+"""
+function divergence!(div, u, setup)
+ (; grid, workgroupsize) = setup
+ (; Δ, N, Ip, Np) = grid
+ D = length(u)
+ e = Offset{D}()
+ @kernel function div!(div, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ d = zero(eltype(div))
+ for α = 1:D
+ d += (u[α][I] - u[α][I-e(α)]) / Δ[α][I[α]]
+ end
+ div[I] = d
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ div!(get_backend(div), workgroupsize)(div, u, I0; ndrange = Np)
+ div
+end
+
+function divergence_adjoint!(u, φ, setup)
+ (; grid, workgroupsize) = setup
+ (; Δ, N, Ip) = grid
+ D = length(u)
+ e = Offset{D}()
+ @kernel function adj!(u, φ)
+ I = @index(Global, Cartesian)
+ for α = 1:D
+ u[α][I] = zero(eltype(u[1]))
+ I ∈ Ip && (u[α][I] += φ[I] / Δ[α][I[α]])
+ I + e(α) ∈ Ip && (u[α][I] -= φ[I+e(α)] / Δ[α][I[α]+1])
+ end
+ end
+ adj!(get_backend(u[1]), workgroupsize)(u, φ; ndrange = N)
+ u
+end
+
+"""
+ divergence(u, setup)
+
+Compute divergence of velocity field.
+"""
+divergence(u, setup) = divergence!(fill!(similar(u[1], setup.grid.N), 0), u, setup)
+
+ChainRulesCore.rrule(::typeof(divergence), u, setup) = (
+ divergence(u, setup),
+ φ -> (
+ NoTangent(),
+ divergence_adjoint!(Tangent{typeof(u)}(similar.(u)...), φ, setup),
+ NoTangent(),
+ ),
+)
+
+"""
+ pressuregradient!(G, p, setup)
+
+Compute pressure gradient (in-place).
+"""
+function pressuregradient!(G, p, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Δu, Nu, Iu) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function G!(G, p, ::Val{α}, I0) where {α}
+ I = @index(Global, Cartesian)
+ I = I0 + I
+ G[α][I] = (p[I+e(α)] - p[I]) / Δu[α][I[α]]
+ end
+ D = dimension()
+ for α = 1:D
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ G!(get_backend(G[1]), workgroupsize)(G, p, Val(α), I0; ndrange = Nu[α])
+ end
+ G
+end
+
+function pressuregradient_adjoint!(pbar, φ, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Δu, N, Iu) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function adj!(p, φ)
+ I = @index(Global, Cartesian)
+ p[I] = zero(eltype(p))
+ for α = 1:D
+ I - e(α) ∈ Iu[α] && (p[I] += φ[α][I-e(α)] / Δu[α][I[α]-1])
+ I ∈ Iu[α] && (p[I] -= φ[α][I] / Δu[α][I[α]])
+ end
+ end
+ adj!(get_backend(pbar), workgroupsize)(pbar, φ; ndrange = N)
+ pbar
+end
+
+"""
+ pressuregradient(p, setup)
+
+Compute pressure gradient.
+"""
+pressuregradient(p, setup) =
+ pressuregradient!(ntuple(α -> zero(p), setup.grid.dimension()), p, setup)
+
+ChainRulesCore.rrule(::typeof(pressuregradient), p, setup) = (
+ pressuregradient(p, setup),
+ φ -> (NoTangent(), pressuregradient_adjoint!(similar(p), (φ...,), setup), NoTangent()),
+)
+
+"""
+ applypressure!(u, p, setup)
+
+Subtract pressure gradient (in-place).
+"""
+function applypressure!(u, p, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Δu, Nu, Iu) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function apply!(u, p, ::Val{α}, I0) where {α}
+ I = @index(Global, Cartesian)
+ I = I0 + I
+ u[α][I] -= (p[I+e(α)] - p[I]) / Δu[α][I[α]]
+ end
+ D = dimension()
+ for α = 1:D
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ apply!(get_backend(u[1]), workgroupsize)(u, p, Val(α), I0; ndrange = Nu[α])
+ end
+ u
+end
+
+# function applypressure_adjoint!(pbar, φ, u, setup)
+# (; grid, workgroupsize) = setup
+# (; dimension, Δu, N, Iu) = grid
+# D = dimension()
+# e = Offset{D}()
+# @kernel function adj!(p, φ)
+# I = @index(Global, Cartesian)
+# p[I] = zero(eltype(p))
+# for α = 1:D
+# I - e(α) ∈ Iu[α] && (p[I] += φ[α][I-e(α)] / Δu[α][I[α]-1])
+# I ∈ Iu[α] && (p[I] -= φ[α][I] / Δu[α][I[α]])
+# end
+# end
+# adj!(get_backend(pbar), workgroupsize)(pbar, φ; ndrange = N)
+# pbar
+# end
+#
+# """
+# applypressure(p, setup)
+#
+# Compute pressure gradient.
+# """
+# applypressure(u, p, setup) =
+# applypressure!(copy.(u), p, setup)
+#
+# ChainRulesCore.rrule(::typeof(applypressure), p, setup) = (
+# applypressure(u, p, setup),
+# φ -> (NoTangent(), applypressure_adjoint!(similar(p), (φ...,), setup), NoTangent()),
+# )
+
+"""
+ laplacian!(L, p, setup)
+
+Compute Laplacian of pressure field (in-place version).
+"""
+function laplacian!(L, p, setup)
+ (; grid, workgroupsize, boundary_conditions) = setup
+ (; dimension, Δ, Δu, N, Np, Ip, Ω) = grid
+ D = dimension()
+ e = Offset{D}()
+ # @kernel function lap!(L, p, I0)
+ # I = @index(Global, Cartesian)
+ # I = I + I0
+ # lap = zero(eltype(p))
+ # for α = 1:D
+ # # bc = boundary_conditions[α]
+ # if bc[1] isa PressureBC && I[α] == I0[α] + 1
+ # lap +=
+ # Ω[I] / Δ[α][I[α]] *
+ # ((p[I+e(α)] - p[I]) / Δu[α][I[α]] - (p[I]) / Δu[α][I[α]-1])
+ # elseif bc[2] isa PressureBC && I[α] == I0[α] + Np[α]
+ # lap +=
+ # Ω[I] / Δ[α][I[α]] *
+ # ((-p[I]) / Δu[α][I[α]] - (p[I] - p[I-e(α)]) / Δu[α][I[α]-1])
+ # elseif bc[1] isa DirichletBC && I[α] == I0[α] + 1
+ # lap += Ω[I] / Δ[α][I[α]] * ((p[I+e(α)] - p[I]) / Δu[α][I[α]])
+ # elseif bc[2] isa DirichletBC && I[α] == I0[α] + Np[α]
+ # lap += Ω[I] / Δ[α][I[α]] * (-(p[I] - p[I-e(α)]) / Δu[α][I[α]-1])
+ # else
+ # lap +=
+ # Ω[I] / Δ[α][I[α]] *
+ # ((p[I+e(α)] - p[I]) / Δu[α][I[α]] - (p[I] - p[I-e(α)]) / Δu[α][I[α]-1])
+ # end
+ # end
+ # L[I] = lap
+ # end
+ @kernel function lapα!(L, p, I0, ::Val{α}, bc) where {α}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ # bc = boundary_conditions[α]
+ if bc[1] isa PressureBC && I[α] == I0[α] + 1
+ L[I] +=
+ Ω[I] / Δ[α][I[α]] *
+ ((p[I+e(α)] - p[I]) / Δu[α][I[α]] - (p[I]) / Δu[α][I[α]-1])
+ elseif bc[2] isa PressureBC && I[α] == I0[α] + Np[α]
+ L[I] +=
+ Ω[I] / Δ[α][I[α]] *
+ ((-p[I]) / Δu[α][I[α]] - (p[I] - p[I-e(α)]) / Δu[α][I[α]-1])
+ elseif bc[1] isa DirichletBC && I[α] == I0[α] + 1
+ L[I] += Ω[I] / Δ[α][I[α]] * ((p[I+e(α)] - p[I]) / Δu[α][I[α]])
+ elseif bc[2] isa DirichletBC && I[α] == I0[α] + Np[α]
+ L[I] += Ω[I] / Δ[α][I[α]] * (-(p[I] - p[I-e(α)]) / Δu[α][I[α]-1])
+ else
+ L[I] +=
+ Ω[I] / Δ[α][I[α]] *
+ ((p[I+e(α)] - p[I]) / Δu[α][I[α]] - (p[I] - p[I-e(α)]) / Δu[α][I[α]-1])
+ end
+ # L[I] = lap
+ end
+ # All volumes have a right velocity
+ # All volumes have a left velocity except the first one
+ # Start at second volume
+ ndrange = Np
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ # lap!(get_backend(L), workgroupsize)(L, p, I0; ndrange)
+ L .= 0
+ for α = 1:D
+ lapα!(get_backend(L), workgroupsize)(
+ L,
+ p,
+ I0,
+ Val(α),
+ boundary_conditions[α];
+ ndrange,
+ )
+ end
+ L
+end
+
+"""
+ laplacian(p, setup)
+
+Compute Laplacian of pressure field.
+"""
+laplacian(p, setup) = laplacian!(similar(p), p, setup)
+
+function laplacian_mat(setup)
+ (; grid, boundary_conditions) = setup
+ (; dimension, x, N, Np, Ip, Δ, Δu, Ω) = grid
+ backend = get_backend(x[1])
+ T = eltype(x[1])
+ D = dimension()
+ e = Offset{D}()
+ Ia = first(Ip)
+ Ib = last(Ip)
+ I = similar(x[1], CartesianIndex{D}, 0)
+ J = similar(x[1], CartesianIndex{D}, 0)
+ val = similar(x[1], 0)
+ I0 = Ia - oneunit(Ia)
+ for α = 1:D
+ a, b = boundary_conditions[α]
+ i = Ip[ntuple(β -> α == β ? (2:Np[α]-1) : (:), D)...][:]
+ ia = Ip[ntuple(β -> α == β ? (1:1) : (:), D)...][:]
+ ib = Ip[ntuple(β -> α == β ? (Np[α]:Np[α]) : (:), D)...][:]
+ for (aa, bb, j) in [(a, nothing, ia), (nothing, nothing, i), (nothing, b, ib)]
+ vala = @.(Ω[j] / Δ[α][getindex.(j, α)] / Δu[α][getindex.(j, α)-1])
+ if isnothing(aa)
+ J = [J; j .- [e(α)]; j]
+ I = [I; j; j]
+ val = [val; vala; -vala]
+ elseif aa isa PressureBC
+ J = [J; j]
+ I = [I; j]
+ val = [val; -vala]
+ elseif aa isa PeriodicBC
+ J = [J; ib; j]
+ I = [I; j; j]
+ val = [val; vala; -vala]
+ elseif aa isa SymmetricBC
+ J = [J; ia; j]
+ I = [I; j; j]
+ val = [val; vala; -vala]
+ elseif aa isa DirichletBC
+ end
+
+ valb = @.(Ω[j] / Δ[α][getindex.(j, α)] / Δu[α][getindex.(j, α)])
+ if isnothing(bb)
+ J = [J; j; j .+ [e(α)]]
+ I = [I; j; j]
+ val = [val; -valb; valb]
+ elseif bb isa PressureBC
+ # The weight of the "right" BC is zero, but still needs a J inside Ip, so
+ # just set it to ib
+ J = [J; j]
+ I = [I; j]
+ val = [val; -valb]
+ elseif bb isa PeriodicBC
+ J = [J; j; ia]
+ I = [I; j; j]
+ val = [val; -valb; valb]
+ elseif bb isa SymmetricBC
+ J = [J; j; ib]
+ I = [I; j; j]
+ val = [val; -valb; valb]
+ elseif bb isa DirichletBC
+ end
+ # val = vcat(
+ # val,
+ # map(I -> Ω[I] / Δ[α][I[α]] / Δu[α][I[α]-1], j),
+ # map(I -> -Ω[I] / Δ[α][I[α]] * (1 / Δu[α][I[α]] + 1 / Δu[α][I[α]-1]), j),
+ # map(I -> Ω[I] / Δ[α][I[α]] / Δu[α][I[α]], j),
+ end
+ end
+ # Go back to CPU, otherwise get following error:
+ # ERROR: CUDA error: an illegal memory access was encountered (code 700, ERROR_ILLEGAL_ADDRESS)
+ I = Array(I)
+ J = Array(J)
+ # I = I .- I0
+ # J = J .- I0
+ I = I .- [I0]
+ J = J .- [I0]
+ # linear = copyto!(similar(x[1], Int, Np), collect(LinearIndices(Ip)))
+ linear = LinearIndices(Ip)
+ I = linear[I]
+ J = linear[J]
+
+ # Assemble on CPU, since CUDA overwrites instead of adding
+ L = sparse(I, J, Array(val))
+ # II = copyto!(similar(x[1], Int, length(I)), I)
+ # JJ = copyto!(similar(x[1], Int, length(J)), J)
+ # sparse(II, JJ, val)
+
+ L
+ # Ω isa CuArray ? cu(L) : L
+end
+
+"""
+ convection!(F, u, setup)
+
+Compute convective term.
+"""
+function convection!(F, u, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Δ, Δu, Nu, Iu, A) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function conv!(F, u, ::Val{α}, ::Val{βrange}, I0) where {α,βrange}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ # for β = 1:D
+ KernelAbstractions.Extras.LoopInfo.@unroll for β in βrange
+ Δuαβ = α == β ? Δu[β] : Δ[β]
+
+ # Half for u[α], (reverse!) interpolation for u[β]
+ # Note:
+ # In matrix version, uses
+ # 1*u[α][I-e(β)] + 0*u[α][I]
+ # instead of 1/2 when u[α][I-e(β)] is at Dirichlet boundary.
+ uαβ1 = (u[α][I-e(β)] + u[α][I]) / 2
+ uαβ2 = (u[α][I] + u[α][I+e(β)]) / 2
+ uβα1 =
+ A[β][α][2][I[α]-(α==β)] * u[β][I-e(β)] +
+ A[β][α][1][I[α]+(α!=β)] * u[β][I-e(β)+e(α)]
+ uβα2 = A[β][α][2][I[α]] * u[β][I] + A[β][α][1][I[α]+1] * u[β][I+e(α)]
+
+ # # Half
+ # uαβ1 = (u[α][I-e(β)] + u[α][I]) / 2
+ # uβα1 = u[β][I-e(β)] / 2 + u[β][I-e(β)+e(α)] / 2
+ # uαβ2 = (u[α][I] + u[α][I+e(β)]) / 2
+ # uβα2 = u[β][I] / 2 + u[β][I+e(α)] / 2
+
+ # # Interpolation
+ # uαβ1 = A[α][β][2][I[β]-1] * u[α][I-e(β)] + A[α][β][1][I[β]] * u[α][I]
+ # uβα1 =
+ # A[β][α][2][I[α]-(α==β)] * u[β][I-e(β)] +
+ # A[β][α][1][I[α]+(α!=β)] * u[β][I-e(β)+e(α)]
+ # uαβ2 = A[α][β][2][I[β]] * u[α][I] + A[α][β][1][I[β]+1] * u[α][I+e(β)]
+ # uβα2 = A[β][α][2][I[α]] * u[β][I] + A[β][α][1][I[α]+1] * u[β][I+e(α)]
+
+ F[α][I] -= (uαβ2 * uβα2 - uαβ1 * uβα1) / Δuαβ[I[β]]
+ end
+ end
+ for α = 1:D
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ conv!(get_backend(F[1]), workgroupsize)(F, u, Val(α), Val(1:D), I0; ndrange = Nu[α])
+ end
+ F
+end
+
+function convection_adjoint!(ubar, φbar, u, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Δ, Δu, N, Iu, A) = grid
+ D = dimension()
+ e = Offset{D}()
+ T = eltype(u[1])
+ h = T(1) / 2
+ @kernel function adj!(ubar, φbar, u, ::Val{γ}, ::Val{looprange}) where {γ,looprange}
+ J = @index(Global, Cartesian)
+ KernelAbstractions.Extras.LoopInfo.@unroll for α in looprange
+ KernelAbstractions.Extras.LoopInfo.@unroll for β in looprange
+ Δuαβ = α == β ? Δu[β] : Δ[β]
+ Aβα1 = A[β][α][1]
+ Aβα2 = A[β][α][2]
+
+ # 1
+ I = J
+ if α == γ && I in Iu[α]
+ uαβ2 = h
+ uβα2 = Aβα2[I[α]] * u[β][I] + Aβα1[I[α]+1] * u[β][I+e(α)]
+ dφdu = -uαβ2 * uβα2 / Δuαβ[I[β]]
+ ubar[γ][J] += φbar[α][I] * dφdu
+ end
+
+ # 2
+ I = J - e(β)
+ if α == γ && I in Iu[α]
+ uαβ2 = h
+ uβα2 = Aβα2[I[α]] * u[β][I] + Aβα1[I[α]+1] * u[β][I+e(α)]
+ dφdu = -uαβ2 * uβα2 / Δuαβ[I[β]]
+ ubar[γ][J] += φbar[α][I] * dφdu
+ end
+
+ # 3
+ I = J
+ if β == γ && I in Iu[α]
+ uαβ2 = h * u[α][I] + h * u[α][I+e(β)]
+ uβα2 = Aβα2[I[α]]
+ dφdu = -uαβ2 * uβα2 / Δuαβ[I[β]]
+ ubar[γ][J] += φbar[α][I] * dφdu
+ end
+
+ # 4
+ I = J - e(α)
+ if β == γ && I in Iu[α]
+ uαβ2 = h * u[α][I] + h * u[α][I+e(β)]
+ uβα2 = Aβα1[I[α]+1]
+ dφdu = -uαβ2 * uβα2 / Δuαβ[I[β]]
+ ubar[γ][J] += φbar[α][I] * dφdu
+ end
+
+ # 5
+ I = J + e(β)
+ if α == γ && I in Iu[α]
+ uαβ1 = h
+ uβα1 =
+ Aβα2[I[α]-(α==β)] * u[β][I-e(β)] +
+ Aβα1[I[α]+(α!=β)] * u[β][I-e(β)+e(α)]
+ dφdu = uαβ1 * uβα1 / Δuαβ[I[β]]
+ ubar[γ][J] += φbar[α][I] * dφdu
+ end
+
+ # 6
+ I = J
+ if α == γ && I in Iu[α]
+ uαβ1 = h
+ uβα1 =
+ Aβα2[I[α]-(α==β)] * u[β][I-e(β)] +
+ Aβα1[I[α]+(α!=β)] * u[β][I-e(β)+e(α)]
+ dφdu = uαβ1 * uβα1 / Δuαβ[I[β]]
+ ubar[γ][J] += φbar[α][I] * dφdu
+ end
+
+ # 7
+ I = J + e(β)
+ if β == γ && I in Iu[α]
+ uαβ1 = h * u[α][I-e(β)] + h * u[α][I]
+ uβα1 = Aβα2[I[α]-(α==β)]
+ dφdu = uαβ1 * uβα1 / Δuαβ[I[β]]
+ ubar[γ][J] += φbar[α][I] * dφdu
+ end
+
+ # 8
+ I = J + e(β) - e(α)
+ if β == γ && I in Iu[α]
+ uαβ1 = h * u[α][I-e(β)] + h * u[α][I]
+ uβα1 = Aβα1[I[α]+(α!=β)]
+ dφdu = uαβ1 * uβα1 / Δuαβ[I[β]]
+ ubar[γ][J] += φbar[α][I] * dφdu
+ end
+ end
+ end
+ end
+ for γ = 1:D
+ adj!(get_backend(u[1]), workgroupsize)(ubar, φbar, u, Val(γ), Val(1:D); ndrange = N)
+ end
+ ubar
+end
+
+convection(u, setup) = convection!(zero.(u), u, setup)
+
+ChainRulesCore.rrule(::typeof(convection), u, setup) = (
+ convection(u, setup),
+ φ -> (
+ NoTangent(),
+ # convection_adjoint!(Tangent{typeof(u)}(zero.(u)...), (φ...,), u, setup),
+ convection_adjoint!(Tangent{typeof(u)}(zero.(u)...), (φ...,), u, setup),
+ NoTangent(),
+ ),
+)
+
+"""
+ diffusion!(F, u, setup)
+
+Compute diffusive term.
+"""
+function diffusion!(F, u, setup)
+ (; grid, workgroupsize, Re) = setup
+ (; dimension, Δ, Δu, Nu, Iu) = grid
+ D = dimension()
+ e = Offset{D}()
+ ν = 1 / Re
+ @kernel function diff!(F, u, ::Val{α}, ::Val{βrange}, I0) where {α,βrange}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ # for β = 1:D
+ KernelAbstractions.Extras.LoopInfo.@unroll for β in βrange
+ Δuαβ = (α == β ? Δu[β] : Δ[β])
+ F[α][I] +=
+ ν * (
+ (u[α][I+e(β)] - u[α][I]) / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) -
+ (u[α][I] - u[α][I-e(β)]) / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1])
+ ) / Δuαβ[I[β]]
+ end
+ end
+ for α = 1:D
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ diff!(get_backend(F[1]), workgroupsize)(F, u, Val(α), Val(1:D), I0; ndrange = Nu[α])
+ end
+ F
+end
+
+function diffusion_adjoint!(u, φ, setup)
+ (; grid, workgroupsize, Re) = setup
+ (; dimension, N, Δ, Δu, Iu) = grid
+ D = dimension()
+ e = Offset{D}()
+ ν = 1 / Re
+ @kernel function adj!(u, φ, ::Val{α}, ::Val{βrange}) where {α,βrange}
+ I = @index(Global, Cartesian)
+ # for β = 1:D
+ KernelAbstractions.Extras.LoopInfo.@unroll for β in βrange
+ Δuαβ = (α == β ? Δu[β] : Δ[β])
+ # F[α][I] += ν * u[α][I+e(β)] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]])
+ # F[α][I] -= ν * u[α][I] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]])
+ # F[α][I] -= ν * u[α][I] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1])
+ # F[α][I] += ν * u[α][I-e(β)] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1])
+ I - e(β) ∈ Iu[α] && (
+ u[α][I] +=
+ ν * φ[α][I-e(β)] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) / Δuαβ[I[β]-1]
+ )
+ I ∈ Iu[α] && (
+ u[α][I] -=
+ ν * φ[α][I] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) / Δuαβ[I[β]]
+ )
+ I ∈ Iu[α] && (
+ u[α][I] -=
+ ν * φ[α][I] / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1]) / Δuαβ[I[β]]
+ )
+ I + e(β) ∈ Iu[α] && (
+ u[α][I] +=
+ ν * φ[α][I+e(β)] / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]]) / Δuαβ[I[β]+1]
+ )
+ end
+ end
+ for α = 1:D
+ adj!(get_backend(u[1]), workgroupsize)(u, φ, Val(α), Val(1:D); ndrange = N)
+ end
+ u
+end
+
+diffusion(u, setup) = diffusion!(zero.(u), u, setup)
+
+ChainRulesCore.rrule(::typeof(diffusion), u, setup) = (
+ diffusion(u, setup),
+ φ -> (
+ NoTangent(),
+ diffusion_adjoint!(Tangent{typeof(u)}(zero.(u)...), (φ...,), setup),
+ NoTangent(),
+ ),
+)
+
+function convectiondiffusion!(F, u, setup)
+ (; grid, workgroupsize, Re) = setup
+ (; dimension, Δ, Δu, Nu, Iu, A) = grid
+ D = dimension()
+ e = Offset{D}()
+ ν = 1 / Re
+ @kernel function cd!(F, u, ::Val{α}, ::Val{βrange}, I0) where {α,βrange}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ # for β = 1:D
+ KernelAbstractions.Extras.LoopInfo.@unroll for β in βrange
+ Δuαβ = α == β ? Δu[β] : Δ[β]
+ uαβ1 = (u[α][I-e(β)] + u[α][I]) / 2
+ uαβ2 = (u[α][I] + u[α][I+e(β)]) / 2
+ uβα1 =
+ A[β][α][2][I[α]-(α==β)] * u[β][I-e(β)] +
+ A[β][α][1][I[α]+(α!=β)] * u[β][I-e(β)+e(α)]
+ uβα2 = A[β][α][2][I[α]] * u[β][I] + A[β][α][1][I[α]+1] * u[β][I+e(α)]
+ uαuβ1 = uαβ1 * uβα1
+ uαuβ2 = uαβ2 * uβα2
+ ∂βuα1 = (u[α][I] - u[α][I-e(β)]) / (β == α ? Δ[β][I[β]] : Δu[β][I[β]-1])
+ ∂βuα2 = (u[α][I+e(β)] - u[α][I]) / (β == α ? Δ[β][I[β]+1] : Δu[β][I[β]])
+ F[α][I] += (ν * (∂βuα2 - ∂βuα1) - (uαuβ2 - uαuβ1)) / Δuαβ[I[β]]
+ end
+ end
+ for α = 1:D
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ cd!(get_backend(F[1]), workgroupsize)(F, u, Val(α), Val(1:D), I0; ndrange = Nu[α])
+ end
+ F
+end
+
+# convectiondiffusion(u, setup) = convectiondiffusion!(zero.(u), u, setup)
+#
+# ChainRulesCore.rrule(::typeof(convectiondiffusion), u, setup) = (
+# convection(u, setup),
+# φ ->
+# (NoTangent(), convectiondiffusion_adjoint!(similar.(u), φ, setup), NoTangent()),
+# )
+
+"""
+ convection_diffusion_temp!(c, u, temp, setup)
+
+Compute convection-diffusion term for the temperature equation.
+Add result to `c`.
+"""
+function convection_diffusion_temp!(c, u, temp, setup)
+ (; grid, workgroupsize, temperature) = setup
+ (; dimension, Δ, Δu, Np, Ip) = grid
+ (; α4) = temperature
+ D = dimension()
+ e = Offset{D}()
+ @kernel function conv!(c, u, temp, ::Val{βrange}, I0) where {βrange}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ cI = zero(eltype(c))
+ KernelAbstractions.Extras.LoopInfo.@unroll for β in βrange
+ # TODO: Add interpolation weights
+ ∂T∂x1 = (temp[I] - temp[I-e(β)]) / Δu[β][I[β]-1]
+ ∂T∂x2 = (temp[I+e(β)] - temp[I]) / Δu[β][I[β]]
+ uT1 = u[β][I-e(β)] * (temp[I] + temp[I-e(β)]) / 2
+ uT2 = u[β][I] * (temp[I+e(β)] + temp[I]) / 2
+ cI += (-(uT2 - uT1) + α4 * (∂T∂x2 - ∂T∂x1)) / Δ[β][I[β]]
+ end
+ c[I] = cI
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ conv!(get_backend(c), workgroupsize)(c, u, temp, Val(1:D), I0; ndrange = Np)
+ c
+end
+
+# function dissipation!(c, u, setup)
+# (; grid, workgroupsize, temperature) = setup
+# (; dimension, Δ, Np, Ip) = grid
+# D = dimension()
+# e = Offset{D}()
+# @inline ∂2(u, α, β, I) = ((u[α][I+e(β)] - u[α][I]) / Δ[β][I])^2 / 2
+# @inline Φ(u, α, β, I) = -∂2(u, α, β, I) - ∂2(u, α, β, I+e(β))
+# @kernel function diss!(d, u, ::Val{βrange}, I0) where {βrange}
+# I = @index(Global, Cartesian)
+# I = I + I0
+# cI = zero(eltype(c))
+# KernelAbstractions.Extras.LoopInfo.@unroll for β in βrange
+# cI += Φ(u, β, β, I) / Δ[β][I[β]]
+# end
+# c[I] += cI
+# end
+# end
+
+"""
+ dissipation!(diss, diff, u, setup)
+
+Compute dissipation term for the temperature equation.
+Add result to `diss`.
+"""
+function dissipation!(diss, diff, u, setup)
+ (; grid, workgroupsize, Re, temperature) = setup
+ (; dimension, Δ, Np, Ip) = grid
+ (; α1, γ) = temperature
+ D = dimension()
+ e = Offset{D}()
+ fill!.(diff, 0)
+ diffusion!(diff, u, setup)
+ @kernel function interpolate!(diss, diff, u, I0, ::Val{βrange}) where {βrange}
+ I = @index(Global, Cartesian)
+ I += I0
+ d = zero(eltype(diss))
+ KernelAbstractions.Extras.LoopInfo.@unroll for β in βrange
+ d += Re * α1 / γ * (u[β][I] * diff[β][I] + u[β][I-e(β)] * diff[β][I-e(β)]) / 2
+ end
+ diss[I] += d
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ interpolate!(get_backend(diss), workgroupsize)(
+ diss,
+ diff,
+ u,
+ I0,
+ Val(1:D);
+ ndrange = Np,
+ )
+ diss
+end
+
+"""
+ bodyforce!(F, u, t, setup)
+
+Compute body force.
+"""
+function bodyforce!(F, u, t, setup)
+ (; grid, workgroupsize, bodyforce, issteadybodyforce) = setup
+ (; dimension, Δ, Δu, Nu, Iu, x, xp) = grid
+ isnothing(bodyforce) && return F
+ D = dimension()
+ e = Offset{D}()
+ @assert D == 2
+ @kernel function f!(F, ::Val{α}, t, I0) where {α}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ # xI = ntuple(β -> α == β ? x[β][1+I[β]] : xp[β][I[β]], D)
+ xI = (
+ α == 1 ? x[1][1+I[1]] : xp[1][I[1]],
+ α == 2 ? x[2][1+I[2]] : xp[2][I[2]],
+ # α == 3 ? x[3][1+I[3]] : xp[3][I[3]],
+ )
+ F[α][I] += bodyforce(Dimension(α), xI..., t)
+ end
+ for α = 1:D
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ if issteadybodyforce
+ F[α] .+= bodyforce[α]
+ else
+ f!(get_backend(F[1]), workgroupsize)(F, Val(α), t, I0; ndrange = Nu[α])
+ end
+ end
+ F
+end
+bodyforce(u, t, setup) = bodyforce!(zero.(u), u, t, setup)
+
+ChainRulesCore.rrule(::typeof(bodyforce), u, t, setup) =
+ (bodyforce(u, t, setup), φ -> (NoTangent(), ZeroTangent(), NoTangent(), NoTangent()))
+
+"""
+ gravity!(F, temp, setup)
+
+Compute gravity term (add to existing `F`).
+"""
+function gravity!(F, temp, setup)
+ (; grid, workgroupsize, temperature) = setup
+ (; dimension, Δ, Δu, Nu, Iu) = grid
+ (; gdir, α2) = temperature
+ D = dimension()
+ e = Offset{D}()
+ @kernel function g!(F, temp, ::Val{gdir}, I0) where {gdir}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ # TODO: Add interpolation weights
+ F[gdir][I] += α2 * (temp[I+e(gdir)] + temp[I]) / 2
+ end
+ I0 = first(Iu[gdir])
+ I0 -= oneunit(I0)
+ g!(get_backend(F[1]), workgroupsize)(F, temp, Val(gdir), I0; ndrange = Nu[gdir])
+ F
+end
+
+"""
+ momentum!(F, u, temp, t, setup)
+
+Right hand side of momentum equations, excluding pressure gradient.
+Put the result in ``F``.
+"""
+function momentum!(F, u, temp, t, setup)
+ (; grid, closure_model, temperature) = setup
+ (; dimension) = grid
+ D = dimension()
+ for α = 1:D
+ F[α] .= 0
+ end
+ # diffusion!(F, u, setup)
+ # convection!(F, u, setup)
+ convectiondiffusion!(F, u, setup)
+ bodyforce!(F, u, t, setup)
+ isnothing(temp) || gravity!(F, temp, setup)
+ F
+end
+
+# monitor(u) = (@info("Forward", typeof(u)); u)
+# ChainRulesCore.rrule(::typeof(monitor), u) =
+# (monitor(u), φ -> (@info("Reverse", typeof(φ)); (NoTangent(), φ)))
+
+# tupleadd(u...) = ntuple(α -> sum(u -> u[α], u), length(u[1]))
+# ChainRulesCore.rrule(::typeof(tupleadd), u...) =
+# (tupleadd(u...), φ -> (NoTangent(), map(u -> φ, u)...))
+
+"""
+ momentum(u, temp, t, setup)
+
+Right hand side of momentum equations, excluding pressure gradient.
+"""
+function momentum(u, temp, t, setup)
+ (; grid, closure_model) = setup
+ (; dimension) = grid
+ D = dimension()
+ d = diffusion(u, setup)
+ c = convection(u, setup)
+ f = bodyforce(u, t, setup)
+ # F = ntuple(D) do α
+ # d[α] .+ c[α] .+ f[α]
+ # end
+ F = @. d + c + f
+ # F = tupleadd(d, c, f)
+ if !isnothing(temp)
+ g = gravity(temp, setup)
+ F = @. F + g
+ end
+ F
+end
+
+# ChainRulesCore.rrule(::typeof(momentum), u, temp, t, setup) = (
+# (error(); momentum(u, temp, t, setup)),
+# φ -> (
+# NoTangent(),
+# momentum_pullback!(zero.(φ), φ, u, temp, t, setup),
+# NoTangent(),
+# NoTangent(),
+# ),
+# )
+
+"""
+ vorticity(u, setup)
+
+Compute vorticity field.
+"""
+vorticity(u, setup) = vorticity!(
+ length(u) == 2 ? similar(u[1], setup.grid.N) :
+ ntuple(α -> similar(u[1], setup.grid.N), length(u)),
+ u,
+ setup,
+)
+
+"""
+ vorticity!(ω, u, setup)
+
+Compute vorticity field.
+"""
+vorticity!(ω, u, setup) = vorticity!(setup.grid.dimension, ω, u, setup)
+
+function vorticity!(::Dimension{2}, ω, u, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Δu, N) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function ω!(ω, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ ω[I] =
+ (u[2][I+e(1)] - u[2][I]) / Δu[1][I[1]] - (u[1][I+e(2)] - u[1][I]) / Δu[2][I[2]]
+ end
+ I0 = CartesianIndex(ntuple(Returns(1), D))
+ I0 -= oneunit(I0)
+ ω!(get_backend(ω), workgroupsize)(ω, u, I0; ndrange = N .- 1)
+ ω
+end
+
+function vorticity!(::Dimension{3}, ω, u, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Δu, N) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function ω!(ω, u, I0)
+ T = eltype(ω)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ for (α, α₊, α₋) in ((1, 2, 3), (2, 3, 1), (3, 1, 2))
+ # α₊ = mod1(α + 1, D)
+ # α₋ = mod1(α - 1, D)
+ ω[α][I] =
+ (u[α₋][I+e(α₊)] - u[α₋][I]) / Δu[α₊][I[α₊]] -
+ (u[α₊][I+e(α₋)] - u[α₊][I]) / Δu[α₋][I[α₋]]
+ end
+ end
+ I0 = CartesianIndex(ntuple(Returns(1), D))
+ I0 -= oneunit(I0)
+ ω!(get_backend(ω[1]), workgroupsize)(ω, u, I0; ndrange = N .- 1)
+ ω
+end
+
+@inline ∂x(uα, I::CartesianIndex{D}, α, β, Δβ, Δuβ; e = Offset{D}()) where {D} =
+ α == β ? (uα[I] - uα[I-e(β)]) / Δβ[I[β]] :
+ (
+ (uα[I+e(β)] - uα[I]) / Δuβ[I[β]] +
+ (uα[I-e(α)+e(β)] - uα[I-e(α)]) / Δuβ[I[β]] +
+ (uα[I] - uα[I-e(β)]) / Δuβ[I[β]-1] +
+ (uα[I-e(α)] - uα[I-e(α)-e(β)]) / Δuβ[I[β]-1]
+ ) / 4
+@inline ∇(u, I::CartesianIndex{2}, Δ, Δu) =
+ @SMatrix [∂x(u[α], I, α, β, Δ[β], Δu[β]) for α = 1:2, β = 1:2]
+@inline ∇(u, I::CartesianIndex{3}, Δ, Δu) =
+ @SMatrix [∂x(u[α], I, α, β, Δ[β], Δu[β]) for α = 1:3, β = 1:3]
+@inline idtensor(u, I::CartesianIndex{2}) =
+ @SMatrix [(α == β) * oneunit(eltype(u[1])) for α = 1:2, β = 1:2]
+@inline idtensor(u, I::CartesianIndex{3}) =
+ @SMatrix [(α == β) * oneunit(eltype(u[1])) for α = 1:3, β = 1:3]
+@inline function strain(u, I, Δ, Δu)
+ ∇u = ∇(u, I, Δ, Δu)
+ (∇u + ∇u') / 2
+end
+@inline gridsize(Δ, I::CartesianIndex{D}) where {D} =
+ sqrt(sum(ntuple(α -> Δ[α][I[α]]^2, D)))
+
+"""
+ smagtensor!(σ, u, θ, setup)
+
+Compute Smagorinsky stress tensors `σ[I]`.
+The Smagorinsky constant `θ` should be a scalar between `0` and `1`.
+"""
+function smagtensor!(σ, u, θ, setup)
+ # TODO: Combine with normal diffusion tensor
+ (; grid, workgroupsize) = setup
+ (; Np, Ip, Δ, Δu) = grid
+ @kernel function σ!(σ, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ S = strain(u, I, Δ, Δu)
+ d = gridsize(Δ, I)
+ νt = θ^2 * d^2 * sqrt(2 * sum(S .* S))
+ σ[I] = 2 * νt * S
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ σ!(get_backend(u[1]), workgroupsize)(σ, u, I0; ndrange = Np)
+ σ
+end
+
+"""
+ divoftensor!(s, σ, setup)
+
+Compute divergence of a tensor with all components in the pressure points.
+The stress tensors should be precomputed and stored in `σ`.
+"""
+function divoftensor!(s, σ, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Nu, Iu, Δ, Δu, A) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function s!(s, σ, ::Val{α}, ::Val{βrange}, I0) where {α,βrange}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ s[α][I] = zero(eltype(s[1]))
+ # for β = 1:D
+ KernelAbstractions.Extras.LoopInfo.@unroll for β in βrange
+ Δuαβ = α == β ? Δu[β] : Δ[β]
+ if α == β
+ σαβ2 = σ[I+e(β)][α, β]
+ σαβ1 = σ[I][α, β]
+ else
+ # TODO: Add interpolation weights for non-uniform case
+ σαβ2 =
+ (
+ σ[I][α, β] +
+ σ[I+e(β)][α, β] +
+ σ[I+e(α)+e(β)][α, β] +
+ σ[I+e(α)][α, β]
+ ) / 4
+ σαβ1 =
+ (
+ σ[I-e(β)][α, β] +
+ σ[I][α, β] +
+ σ[I+e(α)-e(β)][α, β] +
+ σ[I+e(α)][α, β]
+ ) / 4
+ end
+ s[α][I] += (σαβ2 - σαβ1) / Δuαβ[I[β]]
+ end
+ end
+ for α = 1:D
+ I0 = first(Iu[α])
+ I0 -= oneunit(I0)
+ s!(get_backend(s[1]), workgroupsize)(s, σ, Val(α), Val(1:D), I0; ndrange = Nu[α])
+ end
+ s
+end
+
+"""
+ m = smagorinsky_closure(setup)
+
+Create Smagorinsky closure model `m`.
+The model is called as `m(u, θ)`, where the Smagorinsky constant
+`θ` should be a scalar between `0` and `1` (for example `θ = 0.1`).
+"""
+function smagorinsky_closure(setup)
+ (; dimension, x, N) = setup.grid
+ D = dimension()
+ T = eltype(x[1])
+ σ = similar(x[1], SMatrix{D,D,T,D * D}, N)
+ s = ntuple(α -> similar(x[1], N), D)
+ # σ = zero(similar(x[1], SMatrix{D,D,T,D * D}, N))
+ # s = ntuple(α -> zero(similar(x[1], N)), D)
+ function closure(u, θ)
+ smagtensor!(σ, u, θ, setup)
+ apply_bc_p!(σ, zero(T), setup)
+ divoftensor!(s, σ, setup)
+ end
+end
+
+"""
+ tensorbasis!(B, V, u, setup)
+
+Compute symmetry tensor basis `B[1]`-`B[11]` and invariants `V[1]`-`V[5]`,
+as specified in [Silvis2017](@cite) in equations (9) and (11).
+Note that `B[1]` corresponds to ``T_0`` in the paper, and `V` to ``I``.
+"""
+function tensorbasis!(B, V, u, setup)
+ (; grid, workgroupsize) = setup
+ (; Np, Ip, Δ, Δu, dimension) = grid
+ D = dimension()
+ @kernel function basis2!(B, V, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ ∇u = ∇(u, I, Δ, Δu)
+ S = (∇u + ∇u') / 2
+ R = (∇u - ∇u') / 2
+ B[1][I] = idtensor(u, I)
+ B[2][I] = S
+ B[3][I] = S * R - R * S
+ V[1][I] = tr(S * S)
+ V[2][I] = tr(R * R)
+ end
+ @kernel function basis3!(B, V, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ ∇u = ∇(u, I, Δ, Δu)
+ S = (∇u + ∇u') / 2
+ R = (∇u - ∇u') / 2
+ B[1][I] = idtensor(u, I)
+ B[2][I] = S
+ B[3][I] = S * R - R * S
+ B[4][I] = S * S
+ B[5][I] = R * R
+ B[6][I] = S * S * R - R * S * S
+ B[7][I] = S * R * R + R * R * S
+ B[8][I] = R * S * R * R - R * R * S * R
+ B[9][I] = S * R * S * S - S * S * R * S
+ B[10][I] = S * S * R * R + R * R * S * S
+ B[11][I] = R * S * S * R * R - R * R * S * S * R
+ V[1][I] = tr(S * S)
+ V[2][I] = tr(R * R)
+ V[3][I] = tr(S * S * S)
+ V[4][I] = tr(S * R * R)
+ V[5][I] = tr(S * S * R * R)
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ basis! = D == 2 ? basis2! : basis3!
+ basis!(get_backend(u[1]), workgroupsize)(B, V, u, I0; ndrange = Np)
+ B, V
+end
+
+"""
+ tensorbasis(u, setup)
+
+Compute symmetry tensor basis `T[1]`-`T[11]` and invariants `V[1]`-`V[5]`.
+"""
+function tensorbasis(u, setup)
+ T = eltype(u[1])
+ D = setup.grid.dimension()
+ tensorbasis!(
+ ntuple(α -> similar(u[1], SMatrix{D,D,T,D * D}, setup.grid.N), D == 2 ? 3 : 11),
+ ntuple(α -> similar(u[1], setup.grid.N), D == 2 ? 2 : 5),
+ u,
+ setup,
+ )
+end
+
+"""
+ interpolate_u_p(u, setup)
+
+Interpolate velocity to pressure points.
+"""
+interpolate_u_p(u, setup) =
+ interpolate_u_p!(ntuple(α -> similar(u[1], setup.grid.N), length(u)), u, setup)
+
+"""
+ interpolate_u_p!(up, u, setup)
+
+Interpolate velocity to pressure points.
+"""
+function interpolate_u_p!(up, u, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Np, Ip) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function int!(up, u, ::Val{α}, I0) where {α}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ up[α][I] = (u[α][I-e(α)] + u[α][I]) / 2
+ end
+ for α = 1:D
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ int!(get_backend(up[1]), workgroupsize)(up, u, Val(α), I0; ndrange = Np)
+ end
+ up
+end
+
+"""
+ interpolate_ω_p(ω, setup)
+
+Interpolate vorticity to pressure points.
+"""
+interpolate_ω_p(ω, setup) = interpolate_ω_p!(
+ setup.grid.dimension() == 2 ? similar(ω, setup.grid.N) :
+ ntuple(α -> similar(ω[1], setup.grid.N), length(ω)),
+ ω,
+ setup,
+)
+
+"""
+ interpolate_ω_p!(ωp, ω, setup)
+
+Interpolate vorticity to pressure points.
+"""
+interpolate_ω_p!(ωp, ω, setup) = interpolate_ω_p!(setup.grid.dimension, ωp, ω, setup)
+
+function interpolate_ω_p!(::Dimension{2}, ωp, ω, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Np, Ip) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function int!(ωp, ω, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ ωp[I] = (ω[I-e(1)-e(2)] + ω[I]) / 2
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ int!(get_backend(ωp), workgroupsize)(ωp, ω, I0; ndrange = Np)
+ ωp
+end
+
+function interpolate_ω_p!(::Dimension{3}, ωp, ω, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Np, Ip) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function int!(ωp, ω, ::Val{α}, I0) where {α}
+ I = @index(Global, Cartesian)
+ I = I + I0
+ α₊ = mod1(α + 1, D)
+ α₋ = mod1(α - 1, D)
+ ωp[α][I] = (ω[α][I-e(α₊)-e(α₋)] + ω[α][I]) / 2
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ for α = 1:D
+ int!(get_backend(ωp[1]), workgroupsize)(ωp, ω, Val(α), I0; ndrange = Np)
+ end
+ ωp
+end
+
+"""
+ Dfield!(d, G, p, setup; ϵ = eps(eltype(p)))
+
+Compute the ``D``-field [LiJiajia2019](@cite) given by
+
+```math
+D = \\frac{2 | \\nabla p |}{\\nabla^2 p}.
+```
+"""
+function Dfield!(d, G, p, setup; ϵ = eps(eltype(p)))
+ (; grid, workgroupsize) = setup
+ (; dimension, Np, Ip, Δ) = grid
+ T = eltype(p)
+ D = dimension()
+ e = Offset{D}()
+ @kernel function D!(d, G, p, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ g = zero(eltype(p))
+ for α = 1:D
+ g += (G[α][I-e(α)] + G[α][I])^2
+ end
+ lap = zero(eltype(p))
+ # for α = 1:D
+ # lap += (G[α][I] - G[α][I-e(α)]) / Δ[α][I[α]]
+ # end
+ if D == 2
+ lap += (G[1][I] - G[1][I-e(1)]) / Δ[1][I[1]]
+ lap += (G[2][I] - G[2][I-e(2)]) / Δ[2][I[2]]
+ elseif D == 3
+ lap += (G[1][I] - G[1][I-e(1)]) / Δ[1][I[1]]
+ lap += (G[2][I] - G[2][I-e(2)]) / Δ[2][I[2]]
+ lap += (G[3][I] - G[3][I-e(3)]) / Δ[3][I[3]]
+ end
+ lap = lap > 0 ? max(lap, ϵ) : min(lap, -ϵ)
+ # lap = abs(lap)
+ d[I] = sqrt(g) / 2 / lap
+ end
+ pressuregradient!(G, p, setup)
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ D!(get_backend(p), workgroupsize)(d, G, p, I0; ndrange = Np)
+ d
+end
+
+"""
+ Dfield(p, setup; kwargs...)
+
+Compute the ``D``-field.
+"""
+Dfield(p, setup; kwargs...) = Dfield!(
+ zero(p),
+ ntuple(α -> similar(p, setup.grid.N), setup.grid.dimension()),
+ p,
+ setup;
+ kwargs...,
+)
+
+"""
+ Qfield!(Q, u, setup)
+
+Compute ``Q``-field [Jeong1995](@cite) given by
+
+```math
+Q = - \\frac{1}{2} \\sum_{α, β} \\frac{\\partial u^α}{\\partial x^β}
+\\frac{\\partial u^β}{\\partial x^α}.
+```
+"""
+function Qfield!(Q, u, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Np, Ip, Δ) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function Q!(Q, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ q = zero(eltype(Q))
+ for α = 1:D, β = 1:D
+ q -=
+ (u[α][I] - u[α][I-e(β)]) / Δ[β][I[β]] * (u[β][I] - u[β][I-e(α)]) /
+ Δ[α][I[α]] / 2
+ end
+ Q[I] = q
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ Q!(get_backend(u[1]), workgroupsize)(Q, u, I0; ndrange = Np)
+ Q
+end
+
+"""
+ Qfield(u, setup)
+
+Compute the ``Q``-field.
+"""
+Qfield(u, setup) = Qfield!(similar(u[1], setup.grid.N), u, setup)
+
+"""
+ eig2field!(λ, u, setup; ϵ = eps(eltype(λ)))
+
+Compute the second eigenvalue of ``S^2 + \\Omega^2``,
+as proposed by Jeong and Hussain [Jeong1995](@cite).
+"""
+function eig2field!(λ, u, setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Np, Ip, Δ, Δu) = grid
+ D = dimension()
+ @assert D == 3 "eig2 only implemented in 3D"
+ @kernel function λ!(λ, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ ∇u = ∇(u, I, Δ, Δu)
+ S = @. (∇u + ∇u') / 2
+ Ω = @. (∇u - ∇u') / 2
+ # FIXME: Is not recognized as hermitian with Float64 on CPU
+ λ[I] = eigvals(S^2 + Ω^2)[2]
+ end
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ λ!(get_backend(u[1]), workgroupsize)(λ, u, I0; ndrange = Np)
+ λ
+end
+
+"""
+ eig2field(u, setup)
+
+Compute the second eigenvalue of ``S^2 + \\Omega^2``,
+as proposed by Jeong and Hussain [Jeong1995](@cite).
+"""
+eig2field(u, setup) = eig2field!(similar(u[1], setup.grid.N), u, setup)
+
+"""
+ kinetic_energy!(k, u, setup; interpolate_first = false)
+
+Compute kinetic energy field ``k`` (in-place version).
+If `interpolate_first` is true, it is given by
+
+```math
+e_I = \\frac{1}{8} \\sum_\\alpha (u^\\alpha_{I + \\delta(\\alpha) / 2} + u^\\alpha_{I - \\delta(\\alpha) / 2})^2.
+```
+
+Otherwise, it is given by
+
+```math
+e_I = \\frac{1}{4} \\sum_\\alpha (u^\\alpha_{I + \\delta(\\alpha) / 2}^2 + u^\\alpha_{I - \\delta(\\alpha) / 2}^2),
+```
+
+as in [Sanderse2023](@cite).
+"""
+function kinetic_energy!(ke, u, setup; interpolate_first = false)
+ (; grid, workgroupsize) = setup
+ (; dimension, Np, Ip) = grid
+ D = dimension()
+ e = Offset{D}()
+ @kernel function efirst!(ke, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ k = zero(eltype(ke))
+ for α = 1:D
+ k += (u[α][I] + u[α][I-e(α)])^2
+ end
+ k = k / 8
+ ke[I] = k
+ end
+ @kernel function elast!(ke, u, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ k = zero(eltype(ke))
+ for α = 1:D
+ k += u[α][I]^2 + u[α][I-e(α)]^2
+ end
+ k = k / 4
+ ke[I] = k
+ end
+ ke! = interpolate_first ? efirst! : elast!
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ ke!(get_backend(u[1]), workgroupsize)(ke, u, I0; ndrange = Np)
+ ke
+end
+
+"""
+ kinetic_energy(u, setup; kwargs...)
+
+Compute kinetic energy field ``e`` (out-of-place version).
+"""
+kinetic_energy(u, setup; kwargs...) = kinetic_energy!(similar(u[1]), u, setup; kwargs...)
+
+"""
+ total_kinetic_energy(setup, u; kwargs...)
+
+Compute total kinetic energy. The velocity components are interpolated to the
+volume centers and squared.
+"""
+function total_kinetic_energy(u, setup; kwargs...)
+ (; Ω, Ip) = setup.grid
+ k = kinetic_energy(u, setup; kwargs...)
+ k .*= Ω
+ sum(k[Ip])
+end
diff --git a/src/operators/operator_averaging.jl b/src/operators/operator_averaging.jl
deleted file mode 100644
index d39f02c2f..000000000
--- a/src/operators/operator_averaging.jl
+++ /dev/null
@@ -1,411 +0,0 @@
-"""
- operator_averaging(grid, boundary_conditions)
-
-Construct averaging operators.
-"""
-operator_averaging(grid, boundary_conditions) =
- operator_averaging(grid.dimension, grid, boundary_conditions)
-
-# 2D version
-function operator_averaging(::Dimension{2}, grid, boundary_conditions)
- (; Nux_in, Nux_b, Nux_t, Nuy_in, Nuy_b, Nuy_t) = grid
- (; Nvx_in, Nvx_b, Nvx_t, Nvy_in, Nvy_b, Nvy_t) = grid
- (; hx, hy) = grid
- (; order4) = grid
-
- T = eltype(hx)
-
- # Averaging weight:
- weight = T(1 / 2)
-
- ## Averaging operators, u-component
-
- ## Au_ux: evaluate u at ux location
- diag1 = fill(weight, Nux_t - 1)
- A1D = spdiagm(Nux_t - 1, Nux_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Au_ux_bc = bc_general(
- Nux_t,
- Nux_in,
- Nux_b,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Au_ux = I(Nuy_in) ⊗ (A1D * Au_ux_bc.B1D)
- Au_ux_bc = (; Au_ux_bc..., Bbc = I(Nuy_in) ⊗ (A1D * Au_ux_bc.Btemp))
-
- ## Au_uy: evaluate u at uy location
- diag1 = fill(weight, Nuy_t - 1)
- A1D = spdiagm(Nuy_t - 1, Nuy_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Au_uy_bc = bc_general_stag(
- Nuy_t,
- Nuy_in,
- Nuy_b,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Au_uy = (A1D * Au_uy_bc.B1D) ⊗ I(Nux_in)
- Au_uy_bc = (; Au_uy_bc..., Bbc = (A1D * Au_uy_bc.Btemp) ⊗ I(Nux_in))
-
- ## Averaging operators, v-component
-
- ## Av_vx: evaluate v at vx location
- diag1 = fill(weight, Nvx_t - 1)
- A1D = spdiagm(Nvx_t - 1, Nvx_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Av_vx_bc = bc_general_stag(
- Nvx_t,
- Nvx_in,
- Nvx_b,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Av_vx = I(Nvy_in) ⊗ (A1D * Av_vx_bc.B1D)
- Av_vx_bc = (; Av_vx_bc..., Bbc = I(Nvy_in) ⊗ (A1D * Av_vx_bc.Btemp))
-
- ## Av_vy: evaluate v at vy location
- diag1 = fill(weight, Nvy_t - 1)
- A1D = spdiagm(Nvy_t - 1, Nvy_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Av_vy_bc = bc_general(
- Nvy_t,
- Nvy_in,
- Nvy_b,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Av_vy = (A1D * Av_vy_bc.B1D) ⊗ I(Nvx_in)
- Av_vy_bc = (; Av_vy_bc..., Bbc = (A1D * Av_vy_bc.Btemp) ⊗ I(Nvx_in))
-
- ## Fourth order
- if order4
- ## Au_ux: evaluate u at ux location
- diag1 = fill(weight, Nux_in + 3)
- A1D3 = spdiagm(Nux_in + 3, Nux_t + 4, 0 => diag1, 3 => diag1)
-
- # Boundary conditions
- Au_ux_bc3 = bc_av3(
- Nux_t + 4,
- Nux_in,
- Nux_t + 4 - Nux_in,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Au_ux3 = I(Nuy_in) ⊗ (A1D3 * Au_ux_bc3.B1D)
- Au_ux_bc3 = (; Au_ux_bc3..., Bbc = I(Nuy_in) ⊗ (A1D3 * Au_ux_bc3.Btemp))
-
- ## Au_uy: evaluate u at uy location
- diag1 = fill(weight, Nuy_in + 3)
- A1D3 = spdiagm(Nuy_in + 3, Nuy_t + 4, 0 => diag1, 3 => diag1)
-
- # Boundary conditions
- Au_uy_bc3 = bc_av_stag3(
- Nuy_t + 4,
- Nuy_in,
- Nuy_t + 4 - Nuy_in,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Au_uy3 = (A1D3 * Au_uy_bc3.B1D) ⊗ I(Nux_in)
- Au_uy_bc3 = (; Au_uy_bc3..., Bbc = (A1D3 * Au_uy_bc3.Btemp) ⊗ I(Nux_in))
-
- ## Av_vx: evaluate v at vx location
- diag1 = fill(weight, Nvx_in + 3)
- A1D3 = spdiagm(Nvx_in + 3, Nvx_t + 4, 0 => diag1, 3 => diag1)
-
- # Boundary conditions
- Av_vx_bc3 = bc_av_stag3(
- Nvx_t + 4,
- Nvx_in,
- Nvx_t + 4 - Nvx_in,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Av_vx3 = I(Nvy_in) ⊗ (A1D3 * Av_vx_bc3.B1D)
- Av_vx_bc3 = (; Av_vx_bc3..., Bbc = I(Nvy_in) ⊗ (A1D3 * Av_vx_bc3.Btemp))
-
- ## Av_vy: evaluate v at vy location
- diag1 = fill(weight, Nvy_in + 3)
- A1D3 = spdiagm(Nvy_in + 3, Nvy_t + 4, 0 => diag1, 3 => diag1)
-
- # Boundary conditions
- Av_vy_bc3 = bc_av3(
- Nvy_t + 4,
- Nvy_in,
- Nvy_t + 4 - Nvy_in,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Av_vy3 = (A1D3 * Av_vy_bc3.B1D) ⊗ I(Nvx_in)
- Av_vy_bc3 = (; Av_vy_bc3..., Bbc = (A1D3 * Av_vy_bc3.Btemp) ⊗ I(Nvx_in))
- end
-
- ## Group operators
- operators = (; Au_ux, Au_uy, Av_vx, Av_vy, Au_ux_bc, Au_uy_bc, Av_vx_bc, Av_vy_bc)
-
- if order4
- operators = (;
- operators...,
- Au_ux3,
- Au_uy3,
- Av_vx3,
- Av_vy3,
- Au_ux_bc3,
- Au_uy_bc3,
- Av_vx_bc3,
- Av_vy_bc3,
- )
- end
-
- operators
-end
-
-# 3D version
-function operator_averaging(::Dimension{3}, grid, boundary_conditions)
- (; Nux_in, Nux_b, Nux_t, Nuy_in, Nuy_b, Nuy_t, Nuz_in, Nuz_b, Nuz_t) = grid
- (; Nvx_in, Nvx_b, Nvx_t, Nvy_in, Nvy_b, Nvy_t, Nvz_in, Nvz_b, Nvz_t) = grid
- (; Nwx_in, Nwx_b, Nwx_t, Nwy_in, Nwy_b, Nwy_t, Nwz_in, Nwz_b, Nwz_t) = grid
- (; hx, hy, hz) = grid
-
- T = eltype(hx)
-
- # Averaging weight:
- weight = T(1 / 2)
-
- ## Averaging operators, u-component
-
- ## Au_ux: evaluate u at ux location
- diag1 = fill(weight, Nux_t - 1)
- A1D = spdiagm(Nux_t - 1, Nux_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Au_ux_bc = bc_general(
- Nux_t,
- Nux_in,
- Nux_b,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 3D
- Au_ux = I(Nuz_in) ⊗ I(Nuy_in) ⊗ (A1D * Au_ux_bc.B1D)
- Au_ux_bc = (; Au_ux_bc..., Bbc = I(Nuz_in) ⊗ I(Nuy_in) ⊗ (A1D * Au_ux_bc.Btemp))
-
- ## Au_uy: evaluate u at uy location
- diag1 = fill(weight, Nuy_t - 1)
- A1D = spdiagm(Nuy_t - 1, Nuy_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Au_uy_bc = bc_general_stag(
- Nuy_t,
- Nuy_in,
- Nuy_b,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 3D
- Au_uy = I(Nuz_in) ⊗ (A1D * Au_uy_bc.B1D) ⊗ I(Nux_in)
- Au_uy_bc = (; Au_uy_bc..., Bbc = I(Nuz_in) ⊗ (A1D * Au_uy_bc.Btemp) ⊗ I(Nux_in))
-
- ## Au_uz: evaluate u at uz location
- diag1 = fill(weight, Nuz_t - 1)
- A1D = spdiagm(Nuz_t - 1, Nuz_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Au_uz_bc = bc_general_stag(
- Nuz_t,
- Nuz_in,
- Nuz_b,
- boundary_conditions.u.z[1],
- boundary_conditions.u.z[2],
- hz[1],
- hz[end],
- )
-
- # Extend to 3D
- Au_uz = (A1D * Au_uz_bc.B1D) ⊗ I(Nuy_in) ⊗ I(Nux_in)
- Au_uz_bc = (; Au_uz_bc..., Bbc = (A1D * Au_uz_bc.Btemp) ⊗ I(Nuy_in) ⊗ I(Nux_in))
-
- ## Averaging operators, v-component
-
- ## Av_vx: evaluate v at vx location
- diag1 = fill(weight, Nvx_t - 1)
- A1D = spdiagm(Nvx_t - 1, Nvx_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Av_vx_bc = bc_general_stag(
- Nvx_t,
- Nvx_in,
- Nvx_b,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 3D
- Av_vx = I(Nvz_in) ⊗ I(Nvy_in) ⊗ (A1D * Av_vx_bc.B1D)
- Av_vx_bc = (; Av_vx_bc..., Bbc = I(Nvz_in) ⊗ I(Nvy_in) ⊗ (A1D * Av_vx_bc.Btemp))
-
- ## Av_vy: evaluate v at vy location
- diag1 = fill(weight, Nvy_t - 1)
- A1D = spdiagm(Nvy_t - 1, Nvy_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Av_vy_bc = bc_general(
- Nvy_t,
- Nvy_in,
- Nvy_b,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 3D
- Av_vy = I(Nvz_in) ⊗ (A1D * Av_vy_bc.B1D) ⊗ I(Nvx_in)
- Av_vy_bc = (; Av_vy_bc..., Bbc = I(Nvz_in) ⊗ (A1D * Av_vy_bc.Btemp) ⊗ I(Nvx_in))
-
- ## Av_vz: evalvate v at vz location
- diag1 = fill(weight, Nvz_t - 1)
- A1D = spdiagm(Nvz_t - 1, Nvz_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Av_vz_bc = bc_general_stag(
- Nvz_t,
- Nvz_in,
- Nvz_b,
- boundary_conditions.v.z[1],
- boundary_conditions.v.z[2],
- hz[1],
- hz[end],
- )
-
- # Extend to 3D
- Av_vz = (A1D * Av_vz_bc.B1D) ⊗ I(Nvy_in) ⊗ I(Nvx_in)
- Av_vz_bc = (; Av_vz_bc..., Bbc = (A1D * Av_vz_bc.Btemp) ⊗ I(Nvy_in) ⊗ I(Nvx_in))
-
- ## Averaging operators, w-component
-
- ## Aw_wx: evaluate w at wx location
- diag1 = fill(weight, Nwx_t - 1)
- A1D = spdiagm(Nwx_t - 1, Nwx_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Aw_wx_bc = bc_general_stag(
- Nwx_t,
- Nwx_in,
- Nwx_b,
- boundary_conditions.w.x[1],
- boundary_conditions.w.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 3D
- Aw_wx = I(Nwz_in) ⊗ I(Nwy_in) ⊗ (A1D * Aw_wx_bc.B1D)
- Aw_wx_bc = (; Aw_wx_bc..., Bbc = I(Nwz_in) ⊗ I(Nwy_in) ⊗ (A1D * Aw_wx_bc.Btemp))
-
- ## Aw_wy: evaluate w at wy location
- diag1 = fill(weight, Nwy_t - 1)
- A1D = spdiagm(Nwy_t - 1, Nwy_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Aw_wy_bc = bc_general_stag(
- Nwy_t,
- Nwy_in,
- Nwy_b,
- boundary_conditions.w.y[1],
- boundary_conditions.w.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 3D
- Aw_wy = I(Nwz_in) ⊗ (A1D * Aw_wy_bc.B1D) ⊗ I(Nwx_in)
- Aw_wy_bc = (; Aw_wy_bc..., Bbc = I(Nwz_in) ⊗ (A1D * Aw_wy_bc.Btemp) ⊗ I(Nwx_in))
-
- ## Aw_wz: evaluate w at wz location
- diag1 = fill(weight, Nwz_t - 1)
- A1D = spdiagm(Nwz_t - 1, Nwz_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Aw_wz_bc = bc_general(
- Nwz_t,
- Nwz_in,
- Nwz_b,
- boundary_conditions.w.z[1],
- boundary_conditions.w.z[2],
- hz[1],
- hz[end],
- )
-
- # Extend to 3D
- Aw_wz = (A1D * Aw_wz_bc.B1D) ⊗ I(Nwy_in) ⊗ I(Nwx_in)
- Aw_wz_bc =
- (; Aw_wz_bc..., Bbc = kron((A1D * Aw_wz_bc.Btemp), kron(I(Nwy_in), I(Nwx_in))))
-
- ## Group operators
- (;
- Au_ux,
- Au_uy,
- Au_uz,
- Av_vx,
- Av_vy,
- Av_vz,
- Aw_wx,
- Aw_wy,
- Aw_wz,
- Au_ux_bc,
- Au_uy_bc,
- Au_uz_bc,
- Av_vx_bc,
- Av_vy_bc,
- Av_vz_bc,
- Aw_wx_bc,
- Aw_wy_bc,
- Aw_wz_bc,
- )
-end
diff --git a/src/operators/operator_convection_diffusion.jl b/src/operators/operator_convection_diffusion.jl
deleted file mode 100644
index 0fae2e851..000000000
--- a/src/operators/operator_convection_diffusion.jl
+++ /dev/null
@@ -1,1071 +0,0 @@
-"""
- operator_convection_diffusion(grid, boundary_conditions)
-
-Construct convection and diffusion operators.
-"""
-operator_convection_diffusion(grid, boundary_conditions) =
- operator_convection_diffusion(grid.dimension, grid, boundary_conditions)
-
-# 2D version
-function operator_convection_diffusion(::Dimension{2}, grid, boundary_conditions)
- (; Nx, Ny) = grid
- (; Nux_in, Nux_b, Nux_t, Nuy_in, Nuy_b, Nuy_t) = grid
- (; Nvx_in, Nvx_b, Nvx_t, Nvy_in, Nvy_b, Nvy_t) = grid
- (; hx, hy, hxi, hyi, hxd, hyd) = grid
- (; gxi, gyi, gxd, gyd) = grid
- (; Buvy, Bvux) = grid
- (; order4, α) = grid
-
- if order4
- (; hxi3, hyi3, gxi3, gyi3, hxd13, hxd3, hyd13, hyd3) = grid
- (; gxd13, gxd3, gyd13, gyd3) = grid
- (; Ωux, Ωuy, Ωvx, Ωvy) = grid
- (; Ωux1, Ωux3, Ωuy1, Ωuy3, Ωvx1, Ωvx3, Ωvy1, Ωvy3) = grid
- end
-
- T = eltype(hx)
-
- ## Convection (differencing) operator Cu
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nux_t - 2)
- D1D = spdiagm(Nux_t - 2, Nux_t - 1, 0 => -diag1, 1 => diag1)
- Cux = I(Nuy_in) ⊗ D1D
- if !order4
- Dux = Diagonal(hyi) ⊗ D1D
- end
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nuy_t - 2)
- D1D = spdiagm(Nuy_t - 2, Nuy_t - 1, 0 => -diag1, 1 => diag1)
- Cuy = D1D ⊗ I(Nux_in)
- if !order4
- Duy = D1D ⊗ Diagonal(gxi)
- end
-
- # Cu = [Cux Cuy]
- # Du = [Dux Duy]
-
- ## Convection (differencing) operator Cv
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nvx_t - 2)
- D1D = spdiagm(Nvx_t - 2, Nvx_t - 1, 0 => -diag1, 1 => diag1)
- Cvx = I(Nvy_in) ⊗ D1D
- if !order4
- Dvx = Diagonal(gyi) ⊗ D1D
- end
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nvy_t - 2)
- D1D = spdiagm(Nvy_t - 2, Nvy_t - 1, 0 => -diag1, 1 => diag1)
- Cvy = D1D ⊗ I(Nvx_in)
- if !order4
- Dvy = D1D ⊗ Diagonal(hxi)
- end
-
- # Cv = [Cvx Cvy]
- # Dv = [Dvx Dvy]
-
- if order4
- ## Fourth order operators
- ## Convection (differencing) operator Cu
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nux_t - 2)
- D1D = spdiagm(Nux_t - 2, Nux_t + 1, 1 => -diag1, 2 => diag1)
- Dux = Diagonal(hyi) ⊗ D1D
-
- # The "second order" Cux is unchanged
- # The "second order" Dux changes, because we also use the "second
- # Order" flux at "fourth order" ghost points (Dux should have the same
- # Size as Dux3)
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nux_t - 2)
- D1D3 = spdiagm(Nux_t - 2, Nux_t + 1, 0 => -diag1, 3 => diag1)
- Cux3 = I(Ny) ⊗ D1D3
- Dux3 = Diagonal(hyi3) ⊗ D1D3
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nuy_t - 2)
- D1D = spdiagm(Nuy_t - 2, Nuy_t + 1, 1 => -diag1, 2 => diag1)
- Duy = D1D ⊗ Diagonal(gxi)
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nuy_t - 2)
- D1D3 = spdiagm(Nuy_t - 2, Nuy_t + 1, 0 => -diag1, 3 => diag1)
-
- # Uncomment for new BC (functions/new)
- if boundary_conditions.u.y[1] == :dirichlet
- D1D3[1, 1] = 1
- D1D3[1, 2] = -2
- end
- if boundary_conditions.u.y[2] == :dirichlet
- D1D3[end, end-1] = 2
- D1D3[end, end] = -1
- end
- Cuy3 = D1D3 ⊗ I(Nux_in)
- Duy3 = D1D3 ⊗ Diagonal(gxi3)
-
- ## Convection (differencing) operator Cv
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nvx_t - 2)
- D1D = spdiagm(Nvx_t - 2, Nvx_t + 1, 1 => -diag1, 2 => diag1)
- Dvx = Diagonal(gyi) ⊗ D1D
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nvx_t - 2)
- D1D3 = spdiagm(Nvx_t - 2, Nvx_t + 1, 0 => -diag1, 3 => diag1)
-
- # Uncomment for new BC (functions/new)
- if boundary_conditions.v.x[1] == :dirichlet
- D1D3[1, 1] = 1
- D1D3[1, 2] = -2
- end
- if boundary_conditions.v.x[2] == :dirichlet
- D1D3[end, end-1] = 2
- D1D3[end, end] = -1
- end
- Cvx3 = I(Nvy_in) ⊗ D1D3
- Dvx3 = Diagonal(gyi3) ⊗ D1D3
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nvy_t - 2)
- D1D = spdiagm(Nvy_t - 2, Nvy_t + 1, 1 => -diag1, 2 => diag1)
- Dvy = D1D ⊗ Diagonal(hxi)
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nvy_t - 2)
- D1D3 = spdiagm(Nvy_t - 2, Nvy_t + 1, 0 => -diag1, 3 => diag1)
- Cvy3 = D1D3 ⊗ I(Nvx_in)
- Dvy3 = D1D3 ⊗ Diagonal(hxi3)
-
- ## Su_ux: evaluate ux
- diag1 = 1 ./ hxd13
- S1D = spdiagm(Nux_in + 3, Nux_t + 4, 1 => -diag1, 2 => diag1)
-
- # Boundary conditions
- Su_ux_bc = bc_diff3(
- Nux_t + 4,
- Nux_in,
- Nux_t + 4 - Nux_in,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Su_ux = Diagonal(Ωux1) * (I(Ny) ⊗ (S1D * Su_ux_bc.B1D))
- Su_ux_bc = (; Su_ux_bc..., Bbc = Diagonal(Ωux1) * (I(Ny) ⊗ (S1D * Su_ux_bc.Btemp)))
-
- diag1 = 1 ./ hxd3
- S1D3 = spdiagm(Nux_in + 3, Nux_t + 4, 0 => -diag1, 3 => diag1)
-
- # Boundary conditions
- Su_ux_bc3 = bc_diff3(
- Nux_t + 4,
- Nux_in,
- Nux_t + 4 - Nux_in,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Su_ux3 = Diagonal(Ωux3) * (I(Nuy_in) ⊗ (S1D3 * Su_ux_bc3.B1D))
- Su_ux_bc3 =
- (; Su_ux_bc3..., Bbc = Diagonal(Ωux3) * (I(Nuy_in) ⊗ (S1D3 * Su_ux_bc3.Btemp)))
-
- ## Su_uy: evaluate uy
- diag1 = 1 ./ gyd13
- S1D = spdiagm(Nuy_in + 3, Nuy_t + 4, 1 => -diag1, 2 => diag1)
-
- # Boundary conditions
- Su_uy_bc = bc_diff_stag3(
- Nuy_t + 4,
- Nuy_in,
- Nuy_t + 4 - Nuy_in,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Su_uy = Diagonal(Ωuy1) * ((S1D * Su_uy_bc.B1D) ⊗ I(Nux_in))
- Su_uy_bc =
- (; Su_uy_bc..., Bbc = Diagonal(Ωuy1) * ((S1D * Su_uy_bc.Btemp) ⊗ I(Nux_in)))
-
- diag1 = 1 ./ gyd3
- S1D3 = spdiagm(Nuy_in + 3, Nuy_t + 4, 0 => -diag1, 3 => diag1)
-
- # Boundary conditions
- Su_uy_bc3 = bc_diff_stag3(
- Nuy_t + 4,
- Nuy_in,
- Nuy_t + 4 - Nuy_in,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Su_uy3 = Diagonal(Ωuy3) * ((S1D3 * Su_uy_bc3.B1D) ⊗ I(Nux_in))
- Su_uy_bc3 =
- (; Su_uy_bc3..., Bbc = Diagonal(Ωuy3) * ((S1D3 * Su_uy_bc3.Btemp) ⊗ I(Nux_in)))
-
- ## Sv_vx: evaluate vx
- diag1 = 1 ./ gxd13
- S1D = spdiagm(Nvx_in + 3, Nvx_t + 4, 1 => -diag1, 2 => diag1)
-
- # Boundary conditions
- Sv_vx_bc = bc_diff_stag3(
- Nvx_t + 4,
- Nvx_in,
- Nvx_t + 4 - Nvx_in,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Sv_vx = Diagonal(Ωvx1) * (I(Nvy_in) ⊗ (S1D * Sv_vx_bc.B1D))
- Sv_vx_bc =
- (; Sv_vx_bc..., Bbc = Diagonal(Ωvx1) * (I(Nvy_in) ⊗ (S1D * Sv_vx_bc.Btemp)))
-
- diag1 = 1 ./ gxd3
- S1D3 = spdiagm(Nvx_in + 3, Nvx_t + 4, 0 => -diag1, 3 => diag1)
-
- # Boundary conditions
- Sv_vx_bc3 = bc_diff_stag3(
- Nvx_t + 4,
- Nvx_in,
- Nvx_t + 4 - Nvx_in,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
- # Extend to 2D
- Sv_vx3 = Diagonal(Ωvx3) * (I(Nvy_in) ⊗ (S1D3 * Sv_vx_bc3.B1D))
- Sv_vx_bc3 =
- (; Sv_vx_bc3..., Bbc = Diagonal(Ωvx3) * (I(Nvy_in) ⊗ (S1D3 * Sv_vx_bc3.Btemp)))
-
- ## Sv_vy: evaluate vy
- diag1 = 1 ./ hyd13
- S1D = spdiagm(Nvy_in + 3, Nvy_t + 4, 1 => -diag1, 2 => diag1)
-
- # Boundary conditions
- Sv_vy_bc = bc_diff3(
- Nvy_t + 4,
- Nvy_in,
- Nvy_t + 4 - Nvy_in,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Sv_vy = Diagonal(Ωvy1) * ((S1D * Sv_vy_bc.B1D) ⊗ I(Nvx_in))
- Sv_vy_bc =
- (; Sv_vy_bc..., Bbc = Diagonal(Ωvy1) * ((S1D * Sv_vy_bc.Btemp) ⊗ I(Nvx_in)))
-
- diag1 = 1 ./ hyd3
- S1D3 = spdiagm(Nvy_in + 3, Nvy_t + 4, 0 => -diag1, 3 => diag1)
-
- # Boundary conditions
- Sv_vy_bc3 = bc_diff3(
- Nvy_t + 4,
- Nvy_in,
- Nvy_t + 4 - Nvy_in,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- # Extend to 2D
- Sv_vy3 = Diagonal(Ωvy3) * ((S1D3 * Sv_vy_bc3.B1D) ⊗ I(Nvx_in))
- Sv_vy_bc3 =
- (; Sv_vy_bc3..., Bbc = Diagonal(Ωvy3) * ((S1D3 * Sv_vy_bc3.Btemp) ⊗ I(Nvx_in)))
- else
- ## Diffusion operator (stress tensor), u-component: similar to averaging, but with mesh sizes
-
- ## Su_ux: evaluate ux
- diag1 = 1 ./ hxd
- S1D = spdiagm(Nux_t - 1, Nux_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Su_ux_bc = bc_general(
- Nux_t,
- Nux_in,
- Nux_b,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Su_ux = I(Ny) ⊗ (S1D * Su_ux_bc.B1D)
- Su_ux_bc = (; Su_ux_bc..., Bbc = I(Ny) ⊗ (S1D * Su_ux_bc.Btemp))
-
- ## Su_uy: evaluate uy
- diag1 = 1 ./ gyd
- S1D = spdiagm(Nuy_t - 1, Nuy_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Su_uy_bc = bc_diff_stag(
- Nuy_t,
- Nuy_in,
- Nuy_b,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Su_uy = (S1D * Su_uy_bc.B1D) ⊗ I(Nux_in)
- Su_uy_bc = (; Su_uy_bc..., Bbc = (S1D * Su_uy_bc.Btemp) ⊗ I(Nux_in))
-
- ## Sv_uy: evaluate vx at uy; same as Iv_uy except for mesh sizes and -diag diag
- diag1 = 1 ./ gxd
- S1D = spdiagm(Nvx_t - 1, Nvx_t, 0 => -diag1, 1 => diag1)
-
- # The restriction is essentially 1D so it can be directly applied to I1D
- S1D = Bvux * S1D
- S2D = I(Nuy_t - 1) ⊗ S1D
-
- # Boundary conditions low/up
- Nb = Nuy_in + 1 - Nvy_in
- Sv_uy_bc_lu = bc_general(
- Nuy_in + 1,
- Nvy_in,
- Nb,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Sv_uy_bc_lu = (; Sv_uy_bc_lu..., B2D = Sv_uy_bc_lu.B1D ⊗ I(Nvx_in))
- Sv_uy_bc_lu = (; Sv_uy_bc_lu..., Bbc = Sv_uy_bc_lu.Btemp ⊗ I(Nvx_in))
-
- # Boundary conditions left/right
- Sv_uy_bc_lr = bc_general_stag(
- Nvx_t,
- Nvx_in,
- Nvx_b,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Take I2D into left/right operators for convenience
- Sv_uy_bc_lr = (; Sv_uy_bc_lr..., B2D = S2D * (I(Nuy_t - 1) ⊗ Sv_uy_bc_lr.B1D))
- Sv_uy_bc_lr = (; Sv_uy_bc_lr..., Bbc = S2D * (I(Nuy_t - 1) ⊗ Sv_uy_bc_lr.Btemp))
-
- # Resulting operator:
- Sv_uy = Sv_uy_bc_lr.B2D * Sv_uy_bc_lu.B2D
-
- ## Diffusion operator (stress tensor), v-component: similar to averaging!
-
- ## Su_vx: evaluate uy at vx. Same as Iu_vx except for mesh sizes and -diag diag
- diag1 = 1 ./ gyd
- S1D = spdiagm(Nuy_t - 1, Nuy_t, 0 => -diag1, 1 => diag1)
- S1D = Buvy * S1D
- S2D = S1D ⊗ I(Nvx_t - 1)
-
- # Boundary conditions low/up
- Su_vx_bc_lu = bc_general_stag(
- Nuy_t,
- Nuy_in,
- Nuy_b,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
- Su_vx_bc_lu = (; Su_vx_bc_lu..., B2D = S2D * (Su_vx_bc_lu.B1D ⊗ I(Nvx_t - 1)))
- Su_vx_bc_lu = (; Su_vx_bc_lu..., Bbc = S2D * (Su_vx_bc_lu.Btemp ⊗ I(Nvx_t - 1)))
-
- # Boundary conditions left/right
- Nb = Nvx_in + 1 - Nux_in
- Su_vx_bc_lr = bc_general(
- Nvx_in + 1,
- Nux_in,
- Nb,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- Su_vx_bc_lr = (; Su_vx_bc_lr..., B2D = I(Nuy_in) ⊗ Su_vx_bc_lr.B1D)
- Su_vx_bc_lr = (; Su_vx_bc_lr..., Bbc = I(Nuy_in) ⊗ Su_vx_bc_lr.Btemp)
-
- # Resulting operator:
- Su_vx = Su_vx_bc_lu.B2D * Su_vx_bc_lr.B2D
-
- ## Sv_vx: evaluate vx
- diag1 = 1 ./ gxd
- S1D = spdiagm(Nvx_t - 1, Nvx_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Sv_vx_bc = bc_diff_stag(
- Nvx_t,
- Nvx_in,
- Nvx_b,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Sv_vx = I(Nvy_in) ⊗ (S1D * Sv_vx_bc.B1D)
-
- Sv_vx_bc = (; Sv_vx_bc..., Bbc = I(Nvy_in) ⊗ (S1D * Sv_vx_bc.Btemp))
-
- ## Sv_vy: evaluate vy
- diag1 = 1 ./ hyd
- S1D = spdiagm(Nvy_t - 1, Nvy_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Sv_vy_bc = bc_general(
- Nvy_t,
- Nvy_in,
- Nvy_b,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Sv_vy = (S1D * Sv_vy_bc.B1D) ⊗ I(Nx)
- Sv_vy_bc = (; Sv_vy_bc..., Bbc = (S1D * Sv_vy_bc.Btemp) ⊗ I(Nx))
- end
-
- ## Assemble operators
- if order4
- Diffux_div = (α * Dux - Dux3) * Diagonal(1 ./ Ωux)
- Diffuy_div = (α * Duy - Duy3) * Diagonal(1 ./ Ωuy)
- Diffvx_div = (α * Dvx - Dvx3) * Diagonal(1 ./ Ωvx)
- Diffvy_div = (α * Dvy - Dvy3) * Diagonal(1 ./ Ωvy)
- Diffu = Diffux_div * (α * Su_ux - Su_ux3) + Diffuy_div * (α * Su_uy - Su_uy3)
- Diffv = Diffvx_div * (α * Sv_vx - Sv_vx3) + Diffvy_div * (α * Sv_vy - Sv_vy3)
- else
- Diffu = Dux * Su_ux + Duy * Su_uy
- Diffv = Dvx * Sv_vx + Dvy * Sv_vy
- end
- Diff = blockdiag(Diffu, Diffv)
-
- ## Group operators
- operators = (;
- Cux,
- Cuy,
- Cvx,
- Cvy,
- Su_ux,
- Su_uy,
- Sv_vx,
- Sv_vy,
- Su_ux_bc,
- Su_uy_bc,
- Sv_vx_bc,
- Sv_vy_bc,
- Dux,
- Duy,
- Dvx,
- Dvy,
- Diff,
- Sv_uy,
- Su_vx,
- )
-
- if order4
- operators = (;
- operators...,
- Cux3,
- Cuy3,
- Cvx3,
- Cvy3,
- Su_ux_bc3,
- Su_uy_bc3,
- Sv_vx_bc3,
- Sv_vy_bc3,
- Diffux_div,
- Diffuy_div,
- Diffvx_div,
- Diffvy_div,
- )
- else
- operators = (; operators..., Su_vx_bc_lr, Su_vx_bc_lu, Sv_uy_bc_lr, Sv_uy_bc_lu)
- end
-
- operators
-end
-
-# 3D version
-function operator_convection_diffusion(::Dimension{3}, grid, boundary_conditions)
- (; Nx, Ny, Nz) = grid
- (; Nux_in, Nux_b, Nux_t, Nuy_in, Nuy_b, Nuy_t, Nuz_in, Nuz_b, Nuz_t) = grid
- (; Nvx_in, Nvx_b, Nvx_t, Nvy_in, Nvy_b, Nvy_t, Nvz_in, Nvz_b, Nvz_t) = grid
- (; Nwx_in, Nwx_b, Nwx_t, Nwy_in, Nwy_b, Nwy_t, Nwz_in, Nwz_b, Nwz_t) = grid
- (; hx, hy, hz, hxi, hyi, hzi, hxd, hyd, hzd) = grid
- (; gxi, gyi, gzi, gxd, gyd, gzd) = grid
- (; Bvux, Bwux, Buvy, Bwvy, Buwz, Bvwz) = grid
-
- T = eltype(hx)
-
- ## Convection (differencing) operator Cu
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nux_t - 2)
- M1D = spdiagm(Nux_t - 2, Nux_t - 1, 0 => -diag1, 1 => diag1)
- Cux = I(Nz) ⊗ I(Nuy_in) ⊗ M1D
- Dux = Diagonal(hzi) ⊗ Diagonal(hyi) ⊗ M1D
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nuy_t - 2)
- M1D = spdiagm(Nuy_t - 2, Nuy_t - 1, 0 => -diag1, 1 => diag1)
- Cuy = I(Nz) ⊗ M1D ⊗ I(Nux_in)
- Duy = Diagonal(hzi) ⊗ M1D ⊗ Diagonal(gxi)
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nuz_t - 2)
- M1D = spdiagm(Nuz_t - 2, Nuz_t - 1, 0 => -diag1, 1 => diag1)
- Cuz = M1D ⊗ I(Ny) ⊗ I(Nux_in)
- Duz = M1D ⊗ Diagonal(hyi) ⊗ Diagonal(gxi)
-
- # Cu = [Cux Cuy Cuz]
- # Du = [Dux Duy Duz]
-
- ## Convection (differencing) operator Cv
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nvx_t - 2)
- M1D = spdiagm(Nvx_t - 2, Nvx_t - 1, 0 => -diag1, 1 => diag1)
- Cvx = I(Nz) ⊗ I(Nvy_in) ⊗ M1D
- Dvx = Diagonal(hzi) ⊗ Diagonal(gyi) ⊗ M1D
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nvy_t - 2)
- M1D = spdiagm(Nvy_t - 2, Nvy_t - 1, 0 => -diag1, 1 => diag1)
- Cvy = I(Nz) ⊗ M1D ⊗ I(Nx)
- Dvy = Diagonal(hzi) ⊗ M1D ⊗ Diagonal(hxi)
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nvz_t - 2)
- M1D = spdiagm(Nvz_t - 2, Nvz_t - 1, 0 => -diag1, 1 => diag1)
- Cvz = M1D ⊗ I(Nvy_in) ⊗ I(Nx)
- Dvz = M1D ⊗ Diagonal(gyi) ⊗ Diagonal(hxi)
-
- # Cv = [Cvx Cvy Cvz]
- # Dv = [Dvx Dvy Dvz]
-
- ## Convection (differencing) operator Cw
-
- # Calculates difference from pressure points to velocity points
- diag1 = ones(T, Nwx_t - 2)
- M1D = spdiagm(Nwx_t - 2, Nwx_t - 1, 0 => -diag1, 1 => diag1)
- Cwx = I(Nwz_in) ⊗ I(Ny) ⊗ M1D
- Dwx = Diagonal(gzi) ⊗ Diagonal(hyi) ⊗ M1D
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nwy_t - 2)
- M1D = spdiagm(Nwy_t - 2, Nwy_t - 1, 0 => -diag1, 1 => diag1)
- Cwy = I(Nwz_in) ⊗ M1D ⊗ I(Nx)
- Dwy = Diagonal(gzi) ⊗ M1D ⊗ Diagonal(hxi)
-
- # Calculates difference from corner points to velocity points
- diag1 = ones(T, Nwz_t - 2)
- M1D = spdiagm(Nwz_t - 2, Nwz_t - 1, 0 => -diag1, 1 => diag1)
- Cwz = M1D ⊗ I(Ny) ⊗ I(Nx)
- Dwz = M1D ⊗ Diagonal(hyi) ⊗ Diagonal(hxi)
-
- # Cw = [Cwx Cwy Cwz]
- # Dw = [Dwx Dwy Dwz]
-
- ## Diffusion operator (stress tensor), u-component: similar to averaging, but with mesh sizes
-
- ## Su_ux: evaluate ux
- diag1 = 1 ./ hxd
- S1D = spdiagm(Nux_t - 1, Nux_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Su_ux_bc = bc_general(
- Nux_t,
- Nux_in,
- Nux_b,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 3D
- Su_ux = I(Nuz_in) ⊗ I(Nuy_in) ⊗ (S1D * Su_ux_bc.B1D)
- Su_ux_bc = (; Su_ux_bc..., Bbc = I(Nuz_in) ⊗ I(Nuy_in) ⊗ (S1D * Su_ux_bc.Btemp))
-
- ## Su_uy: evaluate uy
- diag1 = 1 ./ gyd
- S1D = spdiagm(Nuy_t - 1, Nuy_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Su_uy_bc = bc_general_stag_diff(
- Nuy_t,
- Nuy_in,
- Nuy_b,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 3D
- Su_uy = I(Nuz_in) ⊗ (S1D * Su_uy_bc.B1D) ⊗ I(Nux_in)
- Su_uy_bc = (; Su_uy_bc..., Bbc = I(Nuz_in) ⊗ (S1D * Su_uy_bc.Btemp) ⊗ I(Nux_in))
-
- ## Su_uz: evaluate uz
- diag1 = 1 ./ gzd
- S1D = spdiagm(Nuz_t - 1, Nuz_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Su_uz_bc = bc_general_stag_diff(
- Nuz_t,
- Nuz_in,
- Nuz_b,
- boundary_conditions.u.z[1],
- boundary_conditions.u.z[2],
- hz[1],
- hz[end],
- )
-
- # Extend to 3D
- Su_uz = (S1D * Su_uz_bc.B1D) ⊗ I(Nuy_in) ⊗ I(Nux_in)
- Su_uz_bc = (; Su_uz_bc..., Bbc = (S1D * Su_uz_bc.Btemp) ⊗ I(Nuy_in) ⊗ I(Nux_in))
-
- ## Sv_uy: evaluate vx at uy; same as Iv_uy except for mesh sizes and -diag diag
- diag1 = 1 ./ gxd
- S1D = spdiagm(Nvx_t - 1, Nvx_t, 0 => -diag1, 1 => diag1)
-
- # The restriction is essentially 1D so it can be directly applied to I1D
- S1D = Bvux * S1D
- S3D = I(Nz) ⊗ I(Nuy_t - 1) ⊗ S1D
-
- # Boundary conditions left/right
- Sv_uy_bc_lr = bc_general_stag(
- Nvx_t,
- Nvx_in,
- Nvx_b,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Take I3D into left/right operators for convenience
- Sv_uy_bc_lr = (; Sv_uy_bc_lr..., B3D = S3D * (I(Nz) ⊗ I(Nuy_t - 1) ⊗ Sv_uy_bc_lr.B1D))
- Sv_uy_bc_lr = (; Sv_uy_bc_lr..., Bbc = S3D * (I(Nz) ⊗ I(Nuy_t - 1) ⊗ Sv_uy_bc_lr.Btemp))
-
- # Boundary conditions low/up
- Nb = Nuy_in + 1 - Nvy_in
- Sv_uy_bc_lu = bc_general(
- Nuy_in + 1,
- Nvy_in,
- Nb,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Sv_uy_bc_lu = (; Sv_uy_bc_lu..., B3D = I(Nz) ⊗ Sv_uy_bc_lu.B1D ⊗ I(Nvx_in))
- Sv_uy_bc_lu = (; Sv_uy_bc_lu..., Bbc = I(Nz) ⊗ Sv_uy_bc_lu.Btemp ⊗ I(Nvx_in))
-
- # Resulting operator:
- Sv_uy = Sv_uy_bc_lr.B3D * Sv_uy_bc_lu.B3D
-
- ## Sw_uz: evaluate wx at uz; same as Iw_uz except for mesh sizes and -diag diag
- diag1 = 1 ./ gxd
- S1D = spdiagm(Nwx_t - 1, Nwx_t, 0 => -diag1, 1 => diag1)
-
- # The restriction is essentially 1D so it can be directly applied to I1D
- S1D = Bwux * S1D
- S3D = I(Nz + 1) ⊗ I(Ny) ⊗ S1D
-
- # Boundary conditions left/right
- Sw_uz_bc_lr = bc_general_stag(
- Nwx_t,
- Nwx_in,
- Nwx_b,
- boundary_conditions.w.x[1],
- boundary_conditions.w.x[2],
- hx[1],
- hx[end],
- )
-
- # Take I3D into left/right operators for convenience
- Sw_uz_bc_lr = (; Sw_uz_bc_lr..., B3D = S3D * (I(Nz + 1) ⊗ I(Ny) ⊗ Sw_uz_bc_lr.B1D))
- Sw_uz_bc_lr = (; Sw_uz_bc_lr..., Bbc = S3D * (I(Nz + 1) ⊗ I(Ny) ⊗ Sw_uz_bc_lr.Btemp))
-
- # Boundary conditions back/front
- Nb = Nuz_in + 1 - Nwz_in
- Sw_uz_bc_bf = bc_general(
- Nuz_in + 1,
- Nwz_in,
- Nb,
- boundary_conditions.w.z[1],
- boundary_conditions.w.z[2],
- hz[1],
- hz[end],
- )
- Sw_uz_bc_bf = (; Sw_uz_bc_bf..., B3D = Sw_uz_bc_bf.B1D ⊗ I(Ny) ⊗ I(Nx))
- Sw_uz_bc_bf = (; Sw_uz_bc_bf..., Bbc = Sw_uz_bc_bf.Btemp ⊗ I(Ny) ⊗ I(Nx))
-
- # Resulting operator:
- Sw_uz = Sw_uz_bc_lr.B3D * Sw_uz_bc_bf.B3D
-
- ## Diffusion operator (stress tensor), v-component: similar to averaging!
-
- ## Su_vx: evaluate uy at vx. Same as Iu_vx except for mesh sizes and -diag diag
- diag1 = 1 ./ gyd
- S1D = spdiagm(Nuy_t - 1, Nuy_t, 0 => -diag1, 1 => diag1)
- S1D = Buvy * S1D
- S3D = I(Nz) ⊗ S1D ⊗ I(Nx + 1)
-
- # Boundary conditions left/right
- Nb = Nvx_in + 1 - Nux_in
- Su_vx_bc_lr = bc_general(
- Nvx_in + 1,
- Nux_in,
- Nb,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
- Su_vx_bc_lr = (; Su_vx_bc_lr..., B3D = I(Nz) ⊗ I(Ny) ⊗ Su_vx_bc_lr.B1D)
- Su_vx_bc_lr = (; Su_vx_bc_lr..., Bbc = I(Nz) ⊗ I(Ny) ⊗ Su_vx_bc_lr.Btemp)
-
- # Boundary conditions low/up
- Su_vx_bc_lu = bc_general_stag(
- Nuy_t,
- Nuy_in,
- Nuy_b,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
- Su_vx_bc_lu = (; Su_vx_bc_lu..., B3D = S3D * (I(Nz) ⊗ Su_vx_bc_lu.B1D ⊗ I(Nx + 1)))
- Su_vx_bc_lu = (; Su_vx_bc_lu..., Bbc = S3D * (I(Nz) ⊗ Su_vx_bc_lu.Btemp ⊗ I(Nx + 1)))
-
- # Resulting operator:
- Su_vx = Su_vx_bc_lu.B3D * Su_vx_bc_lr.B3D
-
- ## Sv_vx: evaluate vx
- diag1 = 1 ./ gxd
- S1D = spdiagm(Nvx_t - 1, Nvx_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Sv_vx_bc = bc_general_stag_diff(
- Nvx_t,
- Nvx_in,
- Nvx_b,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 3D
- Sv_vx = I(Nvz_in) ⊗ I(Nvy_in) ⊗ (S1D * Sv_vx_bc.B1D)
- Sv_vx_bc = (; Sv_vx_bc..., Bbc = I(Nvz_in) ⊗ I(Nvy_in) ⊗ (S1D * Sv_vx_bc.Btemp))
-
- ## Sv_vy: evaluate vy
- diag1 = 1 ./ hyd
- S1D = spdiagm(Nvy_t - 1, Nvy_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Sv_vy_bc = bc_general(
- Nvy_t,
- Nvy_in,
- Nvy_b,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 3D
- Sv_vy = I(Nvz_in) ⊗ (S1D * Sv_vy_bc.B1D) ⊗ I(Nx)
- Sv_vy_bc = (; Sv_vy_bc..., Bbc = I(Nvz_in) ⊗ (S1D * Sv_vy_bc.Btemp) ⊗ I(Nx))
-
- ## Sv_vz: evaluate v at vz location
- diag1 = 1 ./ gzd
- S1D = spdiagm(Nvz_t - 1, Nvz_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Sv_vz_bc = bc_general_stag_diff(
- Nvz_t,
- Nvz_in,
- Nvz_b,
- boundary_conditions.v.z[1],
- boundary_conditions.v.z[2],
- hz[1],
- hz[end],
- )
-
- # Extend to 3D
- Sv_vz = (S1D * Sv_vz_bc.B1D) ⊗ I(Nvy_in) ⊗ I(Nvx_in)
- Sv_vz_bc = (; Sv_vz_bc..., Bbc = (S1D * Sv_vz_bc.Btemp) ⊗ I(Nvy_in) ⊗ I(Nvx_in))
-
- ## Sw_vz: evaluate wy at vz location
- diag1 = 1 ./ gyd
- S1D = spdiagm(Nwy_t - 1, Nwy_t, 0 => -diag1, 1 => diag1)
- S1D = Bwvy * S1D
- S3D = I(Nz + 1) ⊗ S1D ⊗ I(Nx)
-
- # Boundary conditions low/up
- Sw_vz_bc_lu = bc_general_stag(
- Nwy_t,
- Nwy_in,
- Nwy_b,
- boundary_conditions.w.y[1],
- boundary_conditions.w.y[2],
- hy[1],
- hy[end],
- )
- Sw_vz_bc_lu = (; Sw_vz_bc_lu..., B3D = S3D * (I(Nz + 1) ⊗ Sw_vz_bc_lu.B1D ⊗ I(Nx)))
- Sw_vz_bc_lu = (; Sw_vz_bc_lu..., Bbc = S3D * (I(Nz + 1) ⊗ Sw_vz_bc_lu.Btemp ⊗ I(Nx)))
-
- # Boundary conditions back/front
- Nb = Nvz_in + 1 - Nwz_in
- Sw_vz_bc_bf = bc_general(
- Nvz_in + 1,
- Nwz_in,
- Nb,
- boundary_conditions.w.z[1],
- boundary_conditions.w.z[2],
- hz[1],
- hz[end],
- )
- Sw_vz_bc_bf = (; Sw_vz_bc_bf..., B3D = Sw_vz_bc_bf.B1D ⊗ I(Ny) ⊗ I(Nx))
- Sw_vz_bc_bf = (; Sw_vz_bc_bf..., Bbc = Sw_vz_bc_bf.Btemp ⊗ I(Ny) ⊗ I(Nx))
-
- # Resulting operator:
- Sw_vz = Sw_vz_bc_lu.B3D * Sw_vz_bc_bf.B3D
-
- ## Diffusion operator (stress tensor), w-component: similar to averaging, but with mesh sizes
-
- ## Sw_wx: evaluate w at wx location
- diag1 = 1 ./ gxd
- S1D = spdiagm(Nwx_t - 1, Nwx_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Sw_wx_bc = bc_general_stag_diff(
- Nwx_t,
- Nwx_in,
- Nwx_b,
- boundary_conditions.w.x[1],
- boundary_conditions.w.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 3D
- Sw_wx = I(Nwz_in) ⊗ I(Nwy_in) ⊗ (S1D * Sw_wx_bc.B1D)
- Sw_wx_bc = (; Sw_wx_bc..., Bbc = I(Nwz_in) ⊗ I(Nwy_in) ⊗ (S1D * Sw_wx_bc.Btemp))
-
- ## Sw_wy: evaluate w at wy location
- diag1 = 1 ./ gyd
- S1D = spdiagm(Nwy_t - 1, Nwy_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Sw_wy_bc = bc_general_stag_diff(
- Nwy_t,
- Nwy_in,
- Nwy_b,
- boundary_conditions.w.y[1],
- boundary_conditions.w.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 3D
- Sw_wy = I(Nwz_in) ⊗ (S1D * Sw_wy_bc.B1D) ⊗ I(Nwx_in)
- Sw_wy_bc = (; Sw_wy_bc..., Bbc = I(Nwz_in) ⊗ (S1D * Sw_wy_bc.Btemp) ⊗ I(Nwx_in))
-
- ## Sw_wz: evaluate w at wz location
- diag1 = 1 ./ hzd
- S1D = spdiagm(Nwz_t - 1, Nwz_t, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Sw_wz_bc = bc_general(
- Nwz_t,
- Nwz_in,
- Nwz_b,
- boundary_conditions.w.z[1],
- boundary_conditions.w.z[2],
- hz[1],
- hz[end],
- )
-
- # Extend to 3D
- Sw_wz = (S1D * Sw_wz_bc.B1D) ⊗ I(Nwy_in) ⊗ I(Nwx_in)
- Sw_wz_bc = (; Sw_wz_bc..., Bbc = (S1D * Sw_wz_bc.Btemp) ⊗ I(Nwy_in) ⊗ I(Nwx_in))
-
- ## Su_wx: evaluate uz at wx
- diag1 = 1 ./ gzd
- S1D = spdiagm(Nuz_t - 1, Nuz_t, 0 => -diag1, 1 => diag1)
-
- # The restriction is essentially 1D so it can be directly applied to I1D
- S1D = Buwz * S1D
- S3D = S1D ⊗ I(Ny) ⊗ I(Nx + 1)
-
- # Boundary conditions back/front
- Su_wx_bc_bf = bc_general_stag(
- Nuz_t,
- Nuz_in,
- Nuz_b,
- boundary_conditions.u.z[1],
- boundary_conditions.u.z[2],
- hz[1],
- hz[end],
- )
-
- # Take I3D into left/right operators for convenience
- Su_wx_bc_bf = (; Su_wx_bc_bf..., B3D = S3D * (Su_wx_bc_bf.B1D ⊗ I(Ny) ⊗ I(Nx + 1)))
- Su_wx_bc_bf = (; Su_wx_bc_bf..., Bbc = S3D * (Su_wx_bc_bf.Btemp ⊗ I(Ny) ⊗ I(Nx + 1)))
-
- # Boundary conditions left/right
- Nb = Nwx_in + 1 - Nux_in
- Su_wx_bc_lr = bc_general(
- Nwx_in + 1,
- Nux_in,
- Nb,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
- Su_wx_bc_lr = (; Su_wx_bc_lr..., B3D = I(Nz) ⊗ I(Ny) ⊗ Su_wx_bc_lr.B1D)
- Su_wx_bc_lr = (; Su_wx_bc_lr..., Bbc = I(Nz) ⊗ I(Ny) ⊗ Su_wx_bc_lr.Btemp)
-
- # Resulting operator:
- Su_wx = Su_wx_bc_bf.B3D * Su_wx_bc_lr.B3D
-
- ## Sv_wy: evaluate vz at wy
- diag1 = 1 ./ gzd
- S1D = spdiagm(Nvz_t - 1, Nvz_t, 0 => -diag1, 1 => diag1)
-
- # The restriction is essentially 1D so it can be directly applied to I1D
- S1D = Bvwz * S1D
- S3D = S1D ⊗ I(Ny + 1) ⊗ I(Nx)
-
- # Boundary conditions back/front
- Sv_wy_bc_bf = bc_general_stag(
- Nvz_t,
- Nvz_in,
- Nvz_b,
- boundary_conditions.v.z[1],
- boundary_conditions.v.z[2],
- hz[1],
- hz[end],
- )
-
- # Take I3D into left/right operators for convenience
- Sv_wy_bc_bf = (; Sv_wy_bc_bf..., B3D = S3D * (Sv_wy_bc_bf.B1D ⊗ I(Ny + 1) ⊗ I(Nx)))
- Sv_wy_bc_bf = (; Sv_wy_bc_bf..., Bbc = S3D * (Sv_wy_bc_bf.Btemp ⊗ I(Ny + 1) ⊗ I(Nx)))
-
- # Boundary conditions low/up
- Nb = Nwy_in + 1 - Nvy_in
- Sv_wy_bc_lu = bc_general(
- Nwy_in + 1,
- Nvy_in,
- Nb,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Sv_wy_bc_lu = (; Sv_wy_bc_lu..., B3D = I(Nz) ⊗ Sv_wy_bc_lu.B1D ⊗ I(Nx))
- Sv_wy_bc_lu = (; Sv_wy_bc_lu..., Bbc = I(Nz) ⊗ Sv_wy_bc_lu.Btemp ⊗ I(Nx))
-
- # Resulting operator:
- Sv_wy = Sv_wy_bc_bf.B3D * Sv_wy_bc_lu.B3D
-
- ## Assemble operators
- Diffu = Dux * Su_ux + Duy * Su_uy + Duz * Su_uz
- Diffv = Dvx * Sv_vx + Dvy * Sv_vy + Dvz * Sv_vz
- Diffw = Dwx * Sw_wx + Dwy * Sw_wy + Dwz * Sw_wz
- Diff = blockdiag(Diffu, Diffv, Diffw)
-
- ## Group opearators
- (;
- Cux,
- Cuy,
- Cuz,
- Cvx,
- Cvy,
- Cvz,
- Cwx,
- Cwy,
- Cwz,
- Su_ux,
- Su_uy,
- Su_uz,
- Sv_vx,
- Sv_vy,
- Sv_vz,
- Sw_wx,
- Sw_wy,
- Sw_wz,
- Su_ux_bc,
- Su_uy_bc,
- Su_uz_bc,
- Sv_vx_bc,
- Sv_vy_bc,
- Sv_vz_bc,
- Sw_wx_bc,
- Sw_wy_bc,
- Sw_wz_bc,
- Sv_uy_bc_lr,
- Sv_uy_bc_lu,
- Sw_uz_bc_lr,
- Sw_uz_bc_bf,
- Su_vx_bc_lr,
- Su_vx_bc_lu,
- Sw_vz_bc_lu,
- Sw_vz_bc_bf,
- Su_wx_bc_lr,
- Su_wx_bc_bf,
- Sv_wy_bc_lu,
- Sv_wy_bc_bf,
- Dux,
- Duy,
- Duz,
- Dvx,
- Dvy,
- Dvz,
- Dwx,
- Dwy,
- Dwz,
- Diff,
- Sv_uy,
- Su_vx,
- Sw_uz,
- Su_wx,
- Sw_vz,
- Sv_wy,
- )
-end
diff --git a/src/operators/operator_divergence.jl b/src/operators/operator_divergence.jl
deleted file mode 100644
index ae1ddf7fc..000000000
--- a/src/operators/operator_divergence.jl
+++ /dev/null
@@ -1,293 +0,0 @@
-"""
- operator_divergence(grid, boundary_conditions)
-
-Construct divergence and gradient operator.
-"""
-operator_divergence(grid, boundary_conditions) =
- operator_divergence(grid.dimension, grid, boundary_conditions)
-
-# 2D version
-function operator_divergence(::Dimension{2}, grid, boundary_conditions)
- (; Npx, Npy) = grid
- (; Nux_in, Nux_b, Nux_t, Nuy_in) = grid
- (; Nvx_in, Nvy_in, Nvy_b, Nvy_t) = grid
- (; hx, hy) = grid
- (; Ω) = grid
- (; order4, α) = grid
-
- bc = boundary_conditions
-
- if order4
- (; hxi3, hyi3) = grid
- end
-
- T = eltype(hx)
-
- ## Divergence operator M
-
- # Note that the divergence matrix M is not square
- mat_hx = Diagonal(hx)
- mat_hy = Diagonal(hy)
-
- # For fourth order: mat_hx3 is defined in operator_interpolation
-
- ## Mx
- # Building blocks consisting of diagonal matrices where the diagonal is
- # equal to constant per block (hy(block)) and changing for next block to
- # hy(block+1)
- diag1 = ones(T, Nux_t - 1)
- M1D = spdiagm(Nux_t - 1, Nux_t, 0 => -diag1, 1 => diag1)
-
- # We only need derivative at inner pressure points, so we map the resulting
- # boundary matrix (restrict)
- diagpos = 0
- bc.u.x[2] == :pressure && bc.u.x[1] == :pressure && (diagpos = 1)
- bc.u.x[2] != :pressure && bc.u.x[1] == :pressure && (diagpos = 1)
- bc.u.x[2] == :pressure && bc.u.x[1] != :pressure && (diagpos = 0)
- bc.u.x[2] == :periodic && bc.u.x[1] == :periodic && (diagpos = 1)
-
- BMx = spdiagm(Npx, Nux_t - 1, diagpos => ones(T, Npx))
- M1D = BMx * M1D
-
- # Extension to 2D to be used in post-processing files
- Bup = I(Nuy_in) ⊗ BMx
-
- # Boundary conditions
- Mx_bc = bc_general(Nux_t, Nux_in, Nux_b, bc.u.x[1], bc.u.x[2], hx[1], hx[end])
- Mx_bc = (; Mx_bc..., Bbc = mat_hy ⊗ (M1D * Mx_bc.Btemp))
-
- # Extend to 2D
- Mx = mat_hy ⊗ (M1D * Mx_bc.B1D)
-
- if order4
- mat_hy3 = Diagonal(hyi3)
- diag1 = ones(T, Nux_t - 1)
- M1D3 = spdiagm(Nux_t - 1, Nux_t + 2, 0 => -diag1, 3 => diag1)
- M1D3 = BMx * M1D3
- Mx_bc3 = bc_div2(
- Nux_t + 2,
- Nux_in,
- Nux_t + 2 - Nux_in,
- bc.u.x[1],
- bc.u.x[2],
- hx[1],
- hx[end],
- )
- Mx3 = mat_hy3 ⊗ (M1D3 * Mx_bc3.B1D)
- Mx_bc3 = (; Mx_bc3..., Bbc = mat_hy3 ⊗ (M1D3 * Mx_bc3.Btemp))
- end
-
- ## My (same as Mx but reversing indices and kron arguments)
- diag1 = ones(T, Nvy_t - 1)
- M1D = spdiagm(Nvy_t - 1, Nvy_t, 0 => -diag1, 1 => diag1)
-
- # We only need derivative at inner pressure points, so we map the resulting
- # boundary matrix (restriction)
- diagpos = 0
- bc.v.y[2] == :pressure && bc.v.y[1] == :pressure && (diagpos = 1)
- bc.v.y[2] != :pressure && bc.v.y[1] == :pressure && (diagpos = 1)
- bc.v.y[2] == :pressure && bc.v.y[1] != :pressure && (diagpos = 0)
- bc.v.y[2] == :periodic && bc.v.y[1] == :periodic && (diagpos = 1)
-
- BMy = spdiagm(Npy, Nvy_t - 1, diagpos => ones(T, Npy))
- M1D = BMy * M1D
-
- # Extension to 2D to be used in post-processing files
- Bvp = BMy ⊗ I(Nvx_in)
-
- # Boundary conditions
- My_bc = bc_general(Nvy_t, Nvy_in, Nvy_b, bc.v.y[1], bc.v.y[2], hy[1], hy[end])
- My_bc = (; My_bc..., Bbc = (M1D * My_bc.Btemp) ⊗ mat_hx)
-
- # Extend to 2D
- My = (M1D * My_bc.B1D) ⊗ mat_hx
-
- if order4
- mat_hx3 = Diagonal(hxi3)
- diag1 = ones(T, Nvy_t - 1)
- M1D3 = spdiagm(Nvy_t - 1, Nvy_t + 2, 0 => -diag1, 3 => diag1)
- M1D3 = BMy * M1D3
- My_bc3 = bc_div2(
- Nvy_t + 2,
- Nvy_in,
- Nvy_t + 2 - Nvy_in,
- bc.v.y[1],
- bc.v.y[2],
- hy[1],
- hy[end],
- )
- My3 = (M1D3 * My_bc3.B1D) ⊗ mat_hx3
- My_bc3 = (; My_bc3..., Bbc = (M1D3 * My_bc3.Btemp) ⊗ mat_hx3)
- end
-
- ## Resulting divergence matrix
- if order4
- Mx = α * Mx - Mx3
- My = α * My - My3
- end
- M = [Mx My]
-
- ## Gradient operator G
-
- # Like in the continuous case, grad = -div^T
- # Note that this also holds for outflow boundary conditions, if the stress
- # on the ouflow boundary is properly taken into account in y_p (often this
- # stress will be zero)
- G = -sparse(M')
-
- ## Pressure matrix for pressure correction method;
- # Also used to make initial data divergence free or compute additional poisson solve
- # Note that the matrix for the pressure is constant in time.
- # Only the right hand side vector changes, so the pressure matrix can be set up outside the time-stepping-loop.
-
- # Laplace = div grad
- A = M * Diagonal(1 ./ Ω) * G
-
- # Check if all the row sums of the pressure matrix are zero, which
- # should be the case if there are no pressure boundary conditions
- if all(≠(:pressure), (bc.u.x..., bc.v.y...))
- if any(!≈(0; atol = sqrt(eps(T))), sum(A; dims = 2))
- @warn "Pressure matrix: not all rowsums are zero!"
- end
- end
-
- ## Group operators
- operators = (; M, Mx_bc, My_bc, G, Bup, Bvp, A)
-
- if order4
- operators = (; operators..., Mx3, My3, Mx_bc3, My_bc3)
- end
-
- operators
-end
-
-# 3D version
-function operator_divergence(::Dimension{3}, grid, boundary_conditions)
- (; Nux_in, Nux_b, Nux_t, Nuy_in, Nuz_in) = grid
- (; Nvx_in, Nvy_in, Nvy_b, Nvy_t, Nvz_in) = grid
- (; Nwx_in, Nwy_in, Nwz_in, Nwz_b, Nwz_t) = grid
- (; Npx, Npy, Npz) = grid
- (; hx, hy, hz) = grid
- (; Ω) = grid
-
- T = eltype(hx)
-
- bc = boundary_conditions
-
- ## Divergence operator M
-
- # Note that the divergence matrix M is not square
- mat_hx = Diagonal(hx)
- mat_hy = Diagonal(hy)
- mat_hz = Diagonal(hz)
-
- ## Mx
- # Building blocks consisting of diagonal matrices where the diagonal is
- # Equal to constant per block (hy(block)) and changing for next block to
- # Hy(block+1)
- diag1 = ones(T, Nux_t - 1)
- M1D = spdiagm(Nux_t - 1, Nux_t, 0 => -diag1, 1 => diag1)
-
- # We only need derivative at inner pressure points, so we map the resulting
- # Boundary matrix (restrict)
- diagpos = 0
- bc.u.x[2] == :pressure && bc.u.x[1] == :pressure && (diagpos = 1)
- bc.u.x[2] != :pressure && bc.u.x[1] == :pressure && (diagpos = 1)
- bc.u.x[2] == :pressure && bc.u.x[1] != :pressure && (diagpos = 0)
- bc.u.x[2] == :periodic && bc.u.x[1] == :periodic && (diagpos = 1)
- BMx = spdiagm(Npx, Nux_t - 1, diagpos => ones(T, Npx))
- M1D = BMx * M1D
-
- # Extension to 3D to be used in post-processing files
- Bup = I(Nuz_in) ⊗ I(Nuy_in) ⊗ BMx
-
- # Boundary conditions
- Mx_bc = bc_general(Nux_t, Nux_in, Nux_b, bc.u.x[1], bc.u.x[2], hx[1], hx[end])
- Mx_bc = (; Mx_bc..., Bbc = mat_hz ⊗ mat_hy ⊗ (M1D * Mx_bc.Btemp))
-
- # Extend to 3D
- Mx = mat_hz ⊗ mat_hy ⊗ (M1D * Mx_bc.B1D)
-
- ## My
- # Same as Mx but reversing indices and kron arguments
- diag1 = ones(T, Nvy_t - 1)
- M1D = spdiagm(Nvy_t - 1, Nvy_t, 0 => -diag1, 1 => diag1)
-
- # We only need derivative at inner pressure points, so we map the resulting
- # Boundary matrix (restriction)
- diagpos = 0
- bc.v.y[2] == :pressure && bc.v.y[1] == :pressure && (diagpos = 1)
- bc.v.y[2] != :pressure && bc.v.y[1] == :pressure && (diagpos = 1)
- bc.v.y[2] == :pressure && bc.v.y[1] != :pressure && (diagpos = 0)
- bc.v.y[2] == :periodic && bc.v.y[1] == :periodic && (diagpos = 1)
- BMy = spdiagm(Npy, Nvy_t - 1, diagpos => ones(T, Npy))
- M1D = BMy * M1D
-
- # Extension to 3D to be used in post-processing files
- Bvp = I(Nvz_in) ⊗ BMy ⊗ I(Nvx_in)
-
- # Boundary conditions
- My_bc = bc_general(Nvy_t, Nvy_in, Nvy_b, bc.v.y[1], bc.v.y[2], hy[1], hy[end])
- My_bc = (; My_bc..., Bbc = mat_hz ⊗ (M1D * My_bc.Btemp) ⊗ mat_hx)
-
- # Extend to 3D
- My = mat_hz ⊗ (M1D * My_bc.B1D) ⊗ mat_hx
-
- ## Mz
- # Same as Mx but reversing indices and kron arguments
- diag1 = ones(T, Nwz_t - 1)
- M1D = spdiagm(Nwz_t - 1, Nwz_t, 0 => -diag1, 1 => diag1)
-
- # We only need derivative at inner pressure points, so we map the resulting
- # Boundary matrix (restriction)
- diagpos = 0
- bc.w.z[1] == :pressure && bc.w.z[2] == :pressure && (diagpos = 1)
- bc.w.z[1] == :pressure && bc.w.z[2] != :pressure && (diagpos = 1)
- bc.w.z[1] != :pressure && bc.w.z[2] == :pressure && (diagpos = 0)
- bc.w.z[1] == :periodic && bc.w.z[2] == :periodic && (diagpos = 1)
-
- BMz = spdiagm(Npz, Nwz_t - 1, diagpos => ones(T, Npz))
- M1D = BMz * M1D
-
- # Extension to 3D to be used in post-processing files
- Bwp = BMz ⊗ I(Nwy_in) ⊗ I(Nwx_in)
-
- # Boundary conditions
- Mz_bc = bc_general(Nwz_t, Nwz_in, Nwz_b, bc.w.z[1], bc.w.z[2], hz[1], hz[end])
- Mz_bc = (; Mz_bc..., Bbc = (M1D * Mz_bc.Btemp) ⊗ mat_hy ⊗ mat_hx)
-
- # Extend to 3D
- Mz = (M1D * Mz_bc.B1D) ⊗ mat_hy ⊗ mat_hx
-
- ## Resulting divergence matrix
- M = [Mx My Mz]
-
- ## Gradient operator G
-
- # Like in the continuous case, grad = -div^T
- # Note that this also holds for outflow boundary conditions, if the stress
- # on the ouflow boundary is properly taken into account in y_p (often this
- # stress will be zero)
- G = -sparse(M')
-
- ## Pressure matrix for pressure correction method;
- # Also used to make initial data divergence free or compute additional poisson solve
- # if !is_steady(problem)
- # Note that the matrix for the pressure is constant in time.
- # Only the right hand side vector changes, so the pressure matrix can be set up
- # outside the time-stepping-loop.
-
- # Laplace = div grad
- A = M * Diagonal(1 ./ Ω) * G
-
- # Check if all the row sums of the pressure matrix are zero, which
- # should be the case if there are no pressure boundary conditions
- if all(≠(:pressure), (bc.u.x..., bc.v.y..., bc.w.z...))
- if any(!≈(0; atol = sqrt(eps(T))), sum(A; dims = 2))
- @warn "Pressure matrix: not all rowsums are zero!"
- end
- end
-
- ## Group operators
- (; M, Mx_bc, My_bc, Mz_bc, G, Bup, Bvp, Bwp, A)
-end
diff --git a/src/operators/operator_filter.jl b/src/operators/operator_filter.jl
deleted file mode 100644
index 43f6d6dbf..000000000
--- a/src/operators/operator_filter.jl
+++ /dev/null
@@ -1,119 +0,0 @@
-"""
- create_top_hat_p(N, M)
-
-`N` fine points and `M` coarse points in each dimension.
-"""
-function create_top_hat_p(N, M)
- s = N ÷ M
- @assert s * M == N
-
- i = 1:M
- j = reshape(1:M, 1, :)
- ij = @. i + M * (j - 1)
- ij = repeat(ij, 1, 1, s, s)
-
- k = reshape(1:s, 1, 1, :)
- l = reshape(1:s, 1, 1, 1, :)
-
- ijkl = @. s * (i - 1) + k + s * N * (j - 1) + N * (l - 1)
-
- z = fill(1 / s^2, N^2)
-
- sparse(ij[:], ijkl[:], z)
-end
-
-"""
- create_top_hat_u(N, M)
-
-`N` fine points and `M` coarse points in each dimension.
-"""
-function create_top_hat_u(N, M)
- s = N ÷ M
- @assert s * M == N
-
- i = 1:M
- j = reshape(1:M, 1, :)
- ij = @. i + M * (j - 1)
- ij = repeat(ij, 1, 1, 1, s)
-
- k = fill(1, 1, 1, 1)
- l = reshape(1:s, 1, 1, 1, :)
-
- ijkl = @. s * (i - 1) + k + s * N * (j - 1) + N * (l - 1)
-
- z = fill(1 / s, N * M)
-
- sparse(ij[:], ijkl[:], z, M^2, N^2)
-end
-
-"""
- create_top_hat_v(N, M)
-
-`N` fine points and `M` coarse points in each dimension.
-"""
-function create_top_hat_v(N, M)
- s = N ÷ M
- @assert s * M == N
-
- i = 1:M
- j = reshape(1:M, 1, :)
- ij = @. i + M * (j - 1)
- ij = repeat(ij, 1, 1, s, 1)
-
- k = reshape(1:s, 1, 1, :, 1)
- l = fill(1, 1, 1, 1, 1)
-
- ijkl = @. s * (i - 1) + k + s * N * (j - 1) + N * (l - 1)
-
- z = fill(1 / s, N * M)
-
- sparse(ij[:], ijkl[:], z, M^2, N^2)
-end
-
-"""
- create_top_hat_velocity(N, M)
-
-`N` fine points and `M` coarse points in each dimension.
-"""
-function create_top_hat_velocity(N, M)
- Wu = create_top_hat_u(N, M)
- Wv = create_top_hat_v(N, M)
- blockdiag(Wu, Wv)
-end
-
-"""
- operator_filter(grid, boundary_conditions)
-
-Construct filtering operator.
-"""
-operator_filter(grid, boundary_conditions, s) =
- operator_filter(grid.dimension, grid, boundary_conditions, s)
-
-# 2D version
-function operator_filter(::Dimension{2}, grid, boundary_conditions, s)
- (; Nx, Ny, hx, hy) = grid
- N = Nx
- M = N ÷ s
-
- @assert s * M == N
-
- # Requirements
- Δx = first(hx)
- Δy = first(hy)
- @assert all(≈(Δx), hx) && all(≈(Δy), hy) "Filter assumes uniform grid"
- @assert all(
- ==((:periodic, :periodic)),
- (boundary_conditions.u.x, boundary_conditions.v.y),
- ) "Filter assumes periodic boundary conditions"
- @assert Nx == Ny
-
- Kp = create_top_hat_p(N, M)
- Ku = create_top_hat_u(N, M)
- Kv = create_top_hat_v(N, M)
- KV = blockdiag(Ku, Kv)
-
- (; KV, Kp)
-end
-
-# 3D version
-function operator_filter(::Dimension{3}, grid, boundary_conditions) end
diff --git a/src/operators/operator_interpolation.jl b/src/operators/operator_interpolation.jl
deleted file mode 100644
index 9608e9734..000000000
--- a/src/operators/operator_interpolation.jl
+++ /dev/null
@@ -1,857 +0,0 @@
-"""
- operator_interpolation(grid, boundary_conditions)
-
-Construct interpolation operators.
-"""
-operator_interpolation(grid, boundary_conditions) =
- operator_interpolation(grid.dimension, grid, boundary_conditions)
-
-# 2D version
-function operator_interpolation(::Dimension{2}, grid, boundary_conditions)
- (; Nx, Ny) = grid
- (; Nux_in, Nux_b, Nux_t, Nuy_in, Nuy_b, Nuy_t) = grid
- (; Nvx_in, Nvx_b, Nvx_t, Nvy_in, Nvy_b, Nvy_t) = grid
- (; hx, hy, hxi, hyi) = grid
- (; Buvy, Bvux) = grid
- (; order4, β) = grid
-
- if order4
- (; hxi3, hyi3, hx3, hy3) = grid
- end
-
- T = eltype(hx)
-
- weight = T(1 / 2)
-
- mat_hx = Diagonal(hxi)
- mat_hy = Diagonal(hyi)
-
- # Periodic boundary conditions
- if boundary_conditions.u.x == (:periodic, :periodic)
- mat_hx2 = spdiagm(Nx + 2, Nx + 2, [hx[end]; hx; hx[1]])
- else
- mat_hx2 = spdiagm(Nx + 2, Nx + 2, [hx[1]; hx; hx[end]])
- end
-
- if boundary_conditions.v.y == (:periodic, :periodic)
- mat_hy2 = spdiagm(Ny + 2, Ny + 2, [hy[end]; hy; hy[1]])
- else
- mat_hy2 = spdiagm(Ny + 2, Ny + 2, [hy[1]; hy; hy[end]])
- end
-
- ## Interpolation operators, u-component
- if order4
- mat_hx3 = Diagonal(hxi3)
- mat_hy3 = Diagonal(hyi3)
-
- weight1 = 1 // 2 * β
- weight2 = 1 // 2 * (1 - β)
-
- # Periodic boundary conditions
- if boundary_conditions.u.x == (:periodic, :periodic)
- mat_hx2 = spdiagm(Nx + 4, Nx + 4, [hx[end-1]; hx[end]; hx; hx[1]; hx[2]])
- mat_hx4 = spdiagm(Nx + 4, Nx + 4, [hx3[end-1]; hx3[end]; hxi3; hx3[1]; hx3[2]])
- else
- mat_hx2 = spdiagm(Nx + 4, Nx + 4, [hx[2]; hx[1]; hx; hx[end]; hx[end-1]])
- mat_hx4 = spdiagm(
- Nx + 4,
- Nx + 4,
- [
- hx[1] + hx[2] + hx[3]
- 2 * hx[1] + hx[2]
- hxi3
- 2 * hx[end] + hx[end-1]
- hx[end] + hx[end-1] + hx[end-2]
- ],
- )
- end
-
- if boundary_conditions.v.y == (:periodic, :periodic)
- mat_hy2 = spdiagm(Ny + 4, Ny + 4, [hy[end-1]; hy[end]; hy; hy[1]; hy[2]])
- mat_hy4 = spdiagm(Ny + 4, Ny + 4, [hy3[end-1]; hy3[end]; hyi3; hy3[1]; hy3[2]])
- else
- mat_hy2 = spdiagm(Ny + 4, Ny + 4, [hy[2]; hy[1]; hy; hy[end]; hy[end-1]])
- mat_hy4 = spdiagm(
- Ny + 4,
- Ny + 4,
- [
- hy[1] + hy[2] + hy[3]
- 2 * hy[1] + hy[2]
- hyi3
- 2 * hy[end] + hy[end-1]
- hy[end] + hy[end-1] + hy[end-2]
- ],
- )
- end
-
- ## Iu_ux
- diag1 = fill(weight1, Nux_t - 1)
- diag2 = fill(weight2, Nux_t - 1)
- I1D = spdiagm(Nux_t - 1, Nux_t + 2, 0 => diag2, 1 => diag1, 2 => diag1, 3 => diag2)
-
- # Boundary conditions
- Iu_ux_bc = bc_int2(
- Nux_t + 2,
- Nux_in,
- Nux_t + 2 - Nux_in,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Iu_ux = mat_hy ⊗ (I1D * Iu_ux_bc.B1D)
- Iu_ux_bc = (; Iu_ux_bc..., Bbc = mat_hy ⊗ (I1D * Iu_ux_bc.Btemp))
-
- ## Iu_ux3
- diag1 = fill(weight1, Nux_in + 3)
- diag2 = fill(weight2, Nux_in + 3)
- I1D3 =
- spdiagm(Nux_in + 3, Nux_t + 4, 0 => diag2, 1 => diag1, 2 => diag1, 3 => diag2)
-
- # Boundary conditions
- Iu_ux_bc3 = bc_int3(
- Nux_t + 4,
- Nux_in,
- Nux_t + 4 - Nux_in,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Iu_ux3 = mat_hy3 ⊗ (I1D3 * Iu_ux_bc3.B1D)
- Iu_ux_bc3 = (; Iu_ux_bc3..., Bbc = mat_hy3 ⊗ (I1D3 * Iu_ux_bc3.Btemp))
-
- ## Iv_uy
- diag1 = fill(weight1, Nvx_t - 1)
- diag2 = fill(weight2, Nvx_t - 1)
- I1D = spdiagm(Nvx_t - 1, Nvx_t + 2, 0 => diag2, 1 => diag1, 2 => diag1, 3 => diag2)
- any(==(:pressure), boundary_conditions.u.x) &&
- @warn "Possible interpolation bug (see https://github.com/bsanderse/INS2D/commit/b8de84dbe151d6de32928563ca8fa5785cce6318)"
-
- # Restrict to u-points
- # The restriction is essentially 1D so it can be directly applied to I1D
- I1D = Bvux * I1D * mat_hx2
- I2D = I(Nuy_t - 1) ⊗ I1D
-
- # Boundary conditions low/up
- Nb = Nuy_in + 1 - Nvy_in
- Iv_uy_bc_lu = bc_general(
- Nuy_in + 1,
- Nvy_in,
- Nb,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Iv_uy_bc_lu = (; Iv_uy_bc_lu..., B2D = Iv_uy_bc_lu.B1D ⊗ I(Nvx_in))
- Iv_uy_bc_lu = (; Iv_uy_bc_lu..., Bbc = Iv_uy_bc_lu.Btemp ⊗ I(Nvx_in))
-
- # Boundary conditions left/right
- Iv_uy_bc_lr = bc_int_mixed_stag2(
- Nvx_t + 2,
- Nvx_in,
- Nvx_t + 2 - Nvx_in,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Take I2D into left/right operators for convenience
- Iv_uy_bc_lr = (; Iv_uy_bc_lr..., B2D = I2D * (I(Nuy_t - 1) ⊗ Iv_uy_bc_lr.B1D))
- Iv_uy_bc_lr = (; Iv_uy_bc_lr..., Bbc = I2D * (I(Nuy_t - 1) ⊗ Iv_uy_bc_lr.Btemp))
-
- # Resulting operator:
- Iv_uy = Iv_uy_bc_lr.B2D * Iv_uy_bc_lu.B2D
-
- ## Iv_uy3
- diag1 = fill(weight1, Nvx_t - 1)
- diag2 = fill(weight2, Nvx_t - 1)
- I1D = spdiagm(Nvx_t - 1, Nvx_t + 2, 0 => diag2, 1 => diag1, 2 => diag1, 3 => diag2)
-
- # Restrict to u-points
- # The restriction is essentially 1D so it can be directly applied to I1D
- I1D = Bvux * I1D * mat_hx4
- I2D = I(Nuy_t + 1) ⊗ I1D
-
- # Boundary conditions low/up
- Nb = Nuy_in + 3 - Nvy_in
- Iv_uy_bc_lu3 = bc_int_mixed2(
- Nuy_in + 3,
- Nvy_in,
- Nb,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Iv_uy_bc_lu3 = (; Iv_uy_bc_lu3..., B2D = Iv_uy_bc_lu3.B1D ⊗ I(Nvx_in))
- Iv_uy_bc_lu3 = (; Iv_uy_bc_lu3..., Bbc = Iv_uy_bc_lu3.Btemp ⊗ I(Nvx_in))
-
- # Boundary conditions left/right
- Iv_uy_bc_lr3 = bc_int_mixed_stag3(
- Nvx_t + 2,
- Nvx_in,
- Nvx_t + 2 - Nvx_in,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Take I2D into left/right operators for convenience
- Iv_uy_bc_lr3 = (; Iv_uy_bc_lr3..., B2D = I2D * (I(Nuy_t + 1) ⊗ Iv_uy_bc_lr3.B1D))
- Iv_uy_bc_lr3 = (; Iv_uy_bc_lr3..., Bbc = I2D * (I(Nuy_t + 1) ⊗ Iv_uy_bc_lr3.Btemp))
-
- # Resulting operator:
- Iv_uy3 = Iv_uy_bc_lr3.B2D * Iv_uy_bc_lu3.B2D
-
- ## Iu_vx
- diag1 = fill(weight1, Nuy_t - 1)
- diag2 = fill(weight2, Nuy_t - 1)
- I1D = spdiagm(Nuy_t - 1, Nuy_t + 2, 0 => diag2, 1 => diag1, 2 => diag1, 3 => diag2)
- any(==(:pressure), boundary_conditions.v.y) &&
- @warn "Possible interpolation bug (see https://github.com/bsanderse/INS2D/commit/b8de84dbe151d6de32928563ca8fa5785cce6318)"
-
- # Restrict to v-points
- I1D = Buvy * I1D * mat_hy2
- I2D = I1D ⊗ I(Nvx_t - 1)
-
- # Boundary conditions low/up
- Iu_vx_bc_lu = bc_int_mixed_stag2(
- Nuy_t + 2,
- Nuy_in,
- Nuy_t + 2 - Nuy_in,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
- Iu_vx_bc_lu = (; Iu_vx_bc_lu..., B2D = I2D * (Iu_vx_bc_lu.B1D ⊗ I(Nvx_t - 1)))
- Iu_vx_bc_lu = (; Iu_vx_bc_lu..., Bbc = I2D * (Iu_vx_bc_lu.Btemp ⊗ I(Nvx_t - 1)))
-
- # Boundary conditions left/right
- Nb = Nvx_in + 1 - Nux_in
- Iu_vx_bc_lr = bc_general(
- Nvx_in + 1,
- Nux_in,
- Nb,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- Iu_vx_bc_lr = (; Iu_vx_bc_lr..., B2D = I(Nuy_in) ⊗ Iu_vx_bc_lr.B1D)
- Iu_vx_bc_lr = (; Iu_vx_bc_lr..., Bbc = I(Nuy_in) ⊗ Iu_vx_bc_lr.Btemp)
-
- # Resulting operator:
- Iu_vx = Iu_vx_bc_lu.B2D * Iu_vx_bc_lr.B2D
-
- ## Iu_vx3
- diag1 = fill(weight1, Nuy_t - 1)
- diag2 = fill(weight2, Nuy_t - 1)
- I1D = spdiagm(Nuy_t - 1, Nuy_t + 2, 0 => diag2, 1 => diag1, 2 => diag1, 3 => diag2)
-
- # Restrict to v-points
- I1D = Buvy * I1D * mat_hy4
- I2D = I1D ⊗ I(Nvx_t + 1)
-
- # Boundary conditions low/up
- Iu_vx_bc_lu3 = bc_int_mixed_stag3(
- Nuy_t + 2,
- Nuy_in,
- Nuy_t + 2 - Nuy_in,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
- Iu_vx_bc_lu3 = (; Iu_vx_bc_lu3..., B2D = I2D * (Iu_vx_bc_lu3.B1D ⊗ I(Nvx_t + 1)))
- Iu_vx_bc_lu3 = (; Iu_vx_bc_lu3..., Bbc = I2D * (Iu_vx_bc_lu3.Btemp ⊗ I(Nvx_t + 1)))
-
- # Boundary conditions left/right
- Nb = Nvx_in + 3 - Nux_in
- Iu_vx_bc_lr3 = bc_int_mixed2(
- Nvx_in + 3,
- Nux_in,
- Nb,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- Iu_vx_bc_lr3 = (; Iu_vx_bc_lr3..., B2D = I(Nuy_in) ⊗ Iu_vx_bc_lr3.B1D)
- Iu_vx_bc_lr3 = (; Iu_vx_bc_lr3..., Bbc = I(Nuy_in) ⊗ Iu_vx_bc_lr3.Btemp)
-
- # Resulting operator:
- Iu_vx3 = Iu_vx_bc_lu3.B2D * Iu_vx_bc_lr3.B2D
-
- ## Iv_vy
- diag1 = fill(weight1, Nvy_t - 1)
- diag2 = fill(weight2, Nvy_t - 1)
- I1D = spdiagm(Nvy_t - 1, Nvy_t + 2, 0 => diag2, 1 => diag1, 2 => diag1, 3 => diag2)
-
- # Boundary conditions
- Iv_vy_bc = bc_int2(
- Nvy_t + 2,
- Nvy_in,
- Nvy_t + 2 - Nvy_in,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Iv_vy = (I1D * Iv_vy_bc.B1D) ⊗ mat_hx
- Iv_vy_bc = (; Iv_vy_bc..., Bbc = (I1D * Iv_vy_bc.Btemp) ⊗ mat_hx)
-
- ## Iv_vy3
- diag1 = fill(weight1, Nvy_in + 3)
- diag2 = fill(weight2, Nvy_in + 3)
- I1D3 =
- spdiagm(Nvy_in + 3, Nvy_t + 4, 0 => diag2, 1 => diag1, 2 => diag1, 3 => diag2)
-
- # Boundary conditions
- Iv_vy_bc3 = bc_int3(
- Nvy_t + 4,
- Nvy_in,
- Nvy_t + 4 - Nvy_in,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Iv_vy3 = (I1D3 * Iv_vy_bc3.B1D) ⊗ mat_hx3
- Iv_vy_bc3 = (; Iv_vy_bc3..., Bbc = (I1D3 * Iv_vy_bc3.Btemp) ⊗ mat_hx3)
- else
- ## Iu_ux
- diag1 = fill(weight, Nux_t - 1)
- I1D = spdiagm(Nux_t - 1, Nux_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Iu_ux_bc = bc_general(
- Nux_t,
- Nux_in,
- Nux_b,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 2D
- Iu_ux = mat_hy ⊗ (I1D * Iu_ux_bc.B1D)
- Iu_ux_bc = (; Iu_ux_bc..., Bbc = mat_hy ⊗ (I1D * Iu_ux_bc.Btemp))
-
- ## Iv_uy
- diag1 = fill(weight, Nvx_t - 1)
- I1D = spdiagm(Nvx_t - 1, Nvx_t, 0 => diag1, 1 => diag1)
- boundary_conditions.u.x[1] == :pressure && (I1D[1, :] ./= 2)
- boundary_conditions.u.x[2] == :pressure && (I1D[end, :] ./= 2)
-
- # The restriction is essentially 1D so it can be directly applied to I1D
- I1D = Bvux * I1D * mat_hx2
- I2D = I(Nuy_t - 1) ⊗ I1D
-
- # Boundary conditions low/up
- Nb = Nuy_in + 1 - Nvy_in
- Iv_uy_bc_lu = bc_general(
- Nuy_in + 1,
- Nvy_in,
- Nb,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Iv_uy_bc_lu = (; Iv_uy_bc_lu..., B2D = Iv_uy_bc_lu.B1D ⊗ I(Nvx_in))
- Iv_uy_bc_lu = (; Iv_uy_bc_lu..., Bbc = Iv_uy_bc_lu.Btemp ⊗ I(Nvx_in))
-
- # Boundary conditions left/right
- Iv_uy_bc_lr = bc_general_stag(
- Nvx_t,
- Nvx_in,
- Nvx_b,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Take I2D into left/right operators for convenience
- Iv_uy_bc_lr = (; Iv_uy_bc_lr..., B2D = I2D * (I(Nuy_t - 1) ⊗ Iv_uy_bc_lr.B1D))
- Iv_uy_bc_lr = (; Iv_uy_bc_lr..., Bbc = I2D * (I(Nuy_t - 1) ⊗ Iv_uy_bc_lr.Btemp))
-
- # Resulting operator:
- Iv_uy = Iv_uy_bc_lr.B2D * Iv_uy_bc_lu.B2D
-
- ## Interpolation operators, v-component
-
- ## Iu_vx
- diag1 = fill(weight, Nuy_t - 1)
- I1D = spdiagm(Nuy_t - 1, Nuy_t, 0 => diag1, 1 => diag1)
- I1D = Buvy * I1D * mat_hy2
- I2D = I1D ⊗ I(Nvx_t - 1)
- boundary_conditions.v.y[1] == :pressure && (I1D[1, :] ./= 2)
- boundary_conditions.v.y[2] == :pressure && (I1D[end, :] ./= 2)
-
- # Boundary conditions low/up
- Iu_vx_bc_lu = bc_general_stag(
- Nuy_t,
- Nuy_in,
- Nuy_b,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
- Iu_vx_bc_lu = (; Iu_vx_bc_lu..., B2D = I2D * (Iu_vx_bc_lu.B1D ⊗ I(Nvx_t - 1)))
- Iu_vx_bc_lu = (; Iu_vx_bc_lu..., Bbc = I2D * (Iu_vx_bc_lu.Btemp ⊗ I(Nvx_t - 1)))
-
- # Boundary conditions left/right
- Nb = Nvx_in + 1 - Nux_in
- Iu_vx_bc_lr = bc_general(
- Nvx_in + 1,
- Nux_in,
- Nb,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- Iu_vx_bc_lr = (; Iu_vx_bc_lr..., B2D = I(Nuy_in) ⊗ Iu_vx_bc_lr.B1D)
- Iu_vx_bc_lr = (; Iu_vx_bc_lr..., Bbc = I(Nuy_in) ⊗ Iu_vx_bc_lr.Btemp)
-
- # Resulting operator:
- Iu_vx = Iu_vx_bc_lu.B2D * Iu_vx_bc_lr.B2D
-
- ## Iv_vy
- diag1 = fill(weight, Nvy_t - 1)
- I1D = spdiagm(Nvy_t - 1, Nvy_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Iv_vy_bc = bc_general(
- Nvy_t,
- Nvy_in,
- Nvy_b,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Iv_vy = (I1D * Iv_vy_bc.B1D) ⊗ mat_hx
- Iv_vy_bc = (; Iv_vy_bc..., Bbc = (I1D * Iv_vy_bc.Btemp) ⊗ mat_hx)
- end
-
- ## Group operators
- operators = (;
- Iu_ux,
- Iv_uy,
- Iu_vx,
- Iv_vy,
- Iu_ux_bc,
- Iv_vy_bc,
- Iv_uy_bc_lr,
- Iv_uy_bc_lu,
- Iu_vx_bc_lr,
- Iu_vx_bc_lu,
- )
-
- if order4
- operators = (;
- operators...,
- Iu_ux3,
- Iv_uy3,
- Iu_vx3,
- Iv_vy3,
- Iu_ux_bc3,
- Iv_uy_bc_lr3,
- Iv_uy_bc_lu3,
- Iu_vx_bc_lr3,
- Iu_vx_bc_lu3,
- Iv_vy_bc3,
- )
- end
-
- operators
-end
-
-# 3D version
-function operator_interpolation(::Dimension{3}, grid, boundary_conditions)
- (; Nx, Ny, Nz) = grid
- (; Nux_in, Nux_b, Nux_t, Nuy_in, Nuy_b, Nuy_t, Nuz_in, Nuz_b, Nuz_t) = grid
- (; Nvx_in, Nvx_b, Nvx_t, Nvy_in, Nvy_b, Nvy_t, Nvz_in, Nvz_b, Nvz_t) = grid
- (; Nwx_in, Nwx_b, Nwx_t, Nwy_in, Nwy_b, Nwy_t, Nwz_in, Nwz_b, Nwz_t) = grid
- (; hx, hy, hz, hxi, hyi, hzi) = grid
- (; Buvy, Buwz, Bvux, Bvwz, Bwux, Bwvy) = grid
-
- T = eltype(hx)
-
- weight = T(1 / 2)
-
- mat_hx = spdiagm(Nx, Nx, hxi)
- mat_hy = spdiagm(Ny, Ny, hyi)
- mat_hz = spdiagm(Nz, Nz, hzi)
-
- # Periodic boundary conditions
- if boundary_conditions.u.x == (:periodic, :periodic)
- mat_hx2 = spdiagm(Nx + 2, Nx + 2, [hx[end]; hx; hx[1]])
- else
- mat_hx2 = spdiagm(Nx + 2, Nx + 2, [hx[1]; hx; hx[end]])
- end
-
- if boundary_conditions.v.y == (:periodic, :periodic)
- mat_hy2 = spdiagm(Ny + 2, Ny + 2, [hy[end]; hy; hy[1]])
- else
- mat_hy2 = spdiagm(Ny + 2, Ny + 2, [hy[1]; hy; hy[end]])
- end
-
- if boundary_conditions.w.z == (:periodic, :periodic)
- mat_hz2 = spdiagm(Nz + 2, Nz + 2, [hz[end]; hz; hz[1]])
- else
- mat_hz2 = spdiagm(Nz + 2, Nz + 2, [hz[1]; hz; hz[end]])
- end
-
- ## Interpolation operators, u-component
-
- ## Iu_ux
- diag1 = fill(weight, Nux_t - 1)
- I1D = spdiagm(Nux_t - 1, Nux_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Iu_ux_bc = bc_general(
- Nux_t,
- Nux_in,
- Nux_b,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend to 3D
- Iu_ux = mat_hz ⊗ mat_hy ⊗ (I1D * Iu_ux_bc.B1D)
- Iu_ux_bc = (; Iu_ux_bc..., Bbc = mat_hz ⊗ mat_hy ⊗ (I1D * Iu_ux_bc.Btemp))
-
- ## Iv_uy
- diag1 = fill(weight, Nvx_t - 1)
- I1D = spdiagm(Nvx_t - 1, Nvx_t, 0 => diag1, 1 => diag1)
- boundary_conditions.u.x[1] == :pressure && (I1D[1, :] ./= 2)
- boundary_conditions.u.x[2] == :pressure && (I1D[end, :] ./= 2)
-
- # The restriction is essentially 1D so it can be directly applied to I1D
- I1D = Bvux * I1D * mat_hx2
- I3D = kron(mat_hz, kron(I(Nuy_t - 1), I1D))
-
- # Boundary conditions low/up
- Nb = Nuy_in + 1 - Nvy_in
- Iv_uy_bc_lu = bc_general(
- Nuy_in + 1,
- Nvy_in,
- Nb,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Iv_uy_bc_lu = (; Iv_uy_bc_lu..., B3D = I(Nz) ⊗ Iv_uy_bc_lu.B1D ⊗ I(Nvx_in))
- Iv_uy_bc_lu = (; Iv_uy_bc_lu..., Bbc = I(Nz) ⊗ Iv_uy_bc_lu.Btemp ⊗ I(Nvx_in))
-
- # Boundary conditions left/right
- Iv_uy_bc_lr = bc_general_stag(
- Nvx_t,
- Nvx_in,
- Nvx_b,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- # Take I2D into left/right operators for convenience
- Iv_uy_bc_lr = (; Iv_uy_bc_lr..., B3D = I3D * (I(Nz) ⊗ I(Nuy_t - 1) ⊗ Iv_uy_bc_lr.B1D))
- Iv_uy_bc_lr = (; Iv_uy_bc_lr..., Bbc = I3D * (I(Nz) ⊗ I(Nuy_t - 1) ⊗ Iv_uy_bc_lr.Btemp))
-
- # Resulting operator:
- Iv_uy = Iv_uy_bc_lr.B3D * Iv_uy_bc_lu.B3D
-
- ## Iw_uz
- diag1 = fill(weight, Nwx_t - 1)
- I1D = spdiagm(Nwx_t - 1, Nwx_t, 0 => diag1, 1 => diag1)
- boundary_conditions.u.x[1] == :pressure && (I1D[1, :] ./= 2)
- boundary_conditions.u.x[2] == :pressure && (I1D[end, :] ./= 2)
-
- I1D = Bwux * I1D * mat_hx2
- I3D = I(Nz + 1) ⊗ mat_hy ⊗ I1D
-
- # Boundary conditions left/right
- Iw_uz_bc_lr = bc_general_stag(
- Nwx_t,
- Nwx_in,
- Nwx_b,
- boundary_conditions.w.x[1],
- boundary_conditions.w.x[2],
- hx[1],
- hx[end],
- )
-
- # Take I3D into left/right operators for convenience
- Iw_uz_bc_lr = (; Iw_uz_bc_lr..., B3D = I3D * (I(Nz + 1) ⊗ I(Ny) ⊗ Iw_uz_bc_lr.B1D))
- Iw_uz_bc_lr = (; Iw_uz_bc_lr..., Bbc = I3D * (I(Nz + 1) ⊗ I(Ny) ⊗ Iw_uz_bc_lr.Btemp))
-
- # Boundary conditions back/front
- Nb = Nuz_in + 1 - Nwz_in
- Iw_uz_bc_bf = bc_general(
- Nuz_in + 1,
- Nwz_in,
- Nb,
- boundary_conditions.w.z[1],
- boundary_conditions.w.z[2],
- hz[1],
- hz[end],
- )
- Iw_uz_bc_bf = (; Iw_uz_bc_bf..., B3D = Iw_uz_bc_bf.B1D ⊗ I(Ny) ⊗ I(Nx))
- Iw_uz_bc_bf = (; Iw_uz_bc_bf..., Bbc = Iw_uz_bc_bf.Btemp ⊗ I(Ny) ⊗ I(Nx))
-
- # Resulting operator:
- Iw_uz = Iw_uz_bc_lr.B3D * Iw_uz_bc_bf.B3D
-
- ## Interpolation operators, v-component
-
- ## Iu_vx
- diag1 = fill(weight, Nuy_t - 1)
- I1D = spdiagm(Nuy_t - 1, Nuy_t, 0 => diag1, 1 => diag1)
- boundary_conditions.v.y[1] == :pressure && (I1D[1, :] ./= 2)
- boundary_conditions.v.y[2] == :pressure && (I1D[end, :] ./= 2)
-
- I1D = Buvy * I1D * mat_hy2
- I3D = mat_hz ⊗ I1D ⊗ I(Nx + 1)
-
- # Boundary conditions low/up
- Iu_vx_bc_lu = bc_general_stag(
- Nuy_t,
- Nuy_in,
- Nuy_b,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
- Iu_vx_bc_lu = (; Iu_vx_bc_lu..., B3D = I3D * (I(Nz) ⊗ Iu_vx_bc_lu.B1D ⊗ I(Nx + 1)))
- Iu_vx_bc_lu = (; Iu_vx_bc_lu..., Bbc = I3D * (I(Nz) ⊗ Iu_vx_bc_lu.Btemp ⊗ I(Nx + 1)))
-
- # Boundary conditions left/right
- Nb = Nvx_in + 1 - Nux_in
- Iu_vx_bc_lr = bc_general(
- Nvx_in + 1,
- Nux_in,
- Nb,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
- Iu_vx_bc_lr = (; Iu_vx_bc_lr..., B3D = I(Nz) ⊗ I(Ny) ⊗ Iu_vx_bc_lr.B1D)
- Iu_vx_bc_lr = (; Iu_vx_bc_lr..., Bbc = I(Nz) ⊗ I(Ny) ⊗ Iu_vx_bc_lr.Btemp)
-
- # Resulting operator:
- Iu_vx = Iu_vx_bc_lu.B3D * Iu_vx_bc_lr.B3D
-
- ## Iv_vy
- diag1 = fill(weight, Nvy_t - 1)
- I1D = spdiagm(Nvy_t - 1, Nvy_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Iv_vy_bc = bc_general(
- Nvy_t,
- Nvy_in,
- Nvy_b,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 3D
- Iv_vy = mat_hz ⊗ (I1D * Iv_vy_bc.B1D) ⊗ mat_hx
- Iv_vy_bc = (; Iv_vy_bc..., Bbc = mat_hz ⊗ (I1D * Iv_vy_bc.Btemp) ⊗ mat_hx)
-
- ## Iw_vz
- diag1 = fill(weight, Nwy_t - 1)
- I1D = spdiagm(Nwy_t - 1, Nwy_t, 0 => diag1, 1 => diag1)
- boundary_conditions.v.y[1] == :pressure && (I1D[1, :] ./= 2)
- boundary_conditions.v.y[2] == :pressure && (I1D[end, :] ./= 2)
-
- I1D = Bwvy * I1D * mat_hy2
- I3D = I(Nz + 1) ⊗ I1D ⊗ mat_hx
-
- # Boundary conditions low/up
- Iw_vz_bc_lu = bc_general_stag(
- Nwy_t,
- Nwy_in,
- Nwy_b,
- boundary_conditions.w.y[1],
- boundary_conditions.w.y[2],
- hy[1],
- hy[end],
- )
- Iw_vz_bc_lu = (; Iw_vz_bc_lu..., B3D = I3D * (I(Nz + 1) ⊗ Iw_vz_bc_lu.B1D ⊗ I(Nx)))
- Iw_vz_bc_lu = (; Iw_vz_bc_lu..., Bbc = I3D * (I(Nz + 1) ⊗ Iw_vz_bc_lu.Btemp ⊗ I(Nx)))
-
- # Boundary conditions left/right
- Nb = Nvz_in + 1 - Nwz_in
- Iw_vz_bc_bf = bc_general(
- Nvz_in + 1,
- Nwz_in,
- Nb,
- boundary_conditions.w.z[1],
- boundary_conditions.w.z[2],
- hz[1],
- hz[end],
- )
- Iw_vz_bc_bf = (; Iw_vz_bc_bf..., B3D = Iw_vz_bc_bf.B1D ⊗ I(Ny) ⊗ I(Nx))
- Iw_vz_bc_bf = (; Iw_vz_bc_bf..., Bbc = Iw_vz_bc_bf.Btemp ⊗ I(Ny) ⊗ I(Nx))
-
- # Resulting operator:
- Iw_vz = Iw_vz_bc_lu.B3D * Iw_vz_bc_bf.B3D
-
- ## Interpolation operators, w-component
-
- ## Iu_wx
- diag1 = fill(weight, Nuz_t - 1)
- I1D = spdiagm(Nuz_t - 1, Nuz_t, 0 => diag1, 1 => diag1)
- boundary_conditions.w.z[1] == :pressure && (I1D[1, :] ./= 2)
- boundary_conditions.w.z[2] == :pressure && (I1D[end, :] ./= 2)
-
- I1D = Buwz * I1D * mat_hz2
- I3D = I1D ⊗ mat_hy ⊗ I(Nx + 1)
-
- # Boundary conditions back/front
- Iu_wx_bc_bf = bc_general_stag(
- Nuz_t,
- Nuz_in,
- Nuz_b,
- boundary_conditions.u.z[1],
- boundary_conditions.u.z[2],
- hz[1],
- hz[end],
- )
- Iu_wx_bc_bf = (; Iu_wx_bc_bf..., B3D = I3D * (Iu_wx_bc_bf.B1D ⊗ I(Ny) ⊗ I(Nx + 1)))
- Iu_wx_bc_bf = (; Iu_wx_bc_bf..., Bbc = I3D * (Iu_wx_bc_bf.Btemp ⊗ I(Ny) ⊗ I(Nx + 1)))
-
- # Boundary conditions left/right
- Nb = Nwx_in + 1 - Nux_in
- Iu_wx_bc_lr = bc_general(
- Nwx_in + 1,
- Nux_in,
- Nb,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
- Iu_wx_bc_lr = (; Iu_wx_bc_lr..., B3D = I(Nz) ⊗ I(Ny) ⊗ Iu_wx_bc_lr.B1D)
- Iu_wx_bc_lr = (; Iu_wx_bc_lr..., Bbc = I(Nz) ⊗ I(Ny) ⊗ Iu_wx_bc_lr.Btemp)
-
- # Resulting operator:
- Iu_wx = Iu_wx_bc_bf.B3D * Iu_wx_bc_lr.B3D
-
- ## Iv_wy
- diag1 = fill(weight, Nvz_t - 1)
- I1D = spdiagm(Nvz_t - 1, Nvz_t, 0 => diag1, 1 => diag1)
- boundary_conditions.w.z[1] == :pressure && (I1D[1, :] ./= 2)
- boundary_conditions.w.z[2] == :pressure && (I1D[end, :] ./= 2)
-
- I1D = Bvwz * I1D * mat_hz2
- I3D = I1D ⊗ I(Ny + 1) ⊗ mat_hx
-
- # Boundary conditions back/front
- Iv_wy_bc_bf = bc_general_stag(
- Nvz_t,
- Nvz_in,
- Nvz_b,
- boundary_conditions.v.z[1],
- boundary_conditions.v.z[2],
- hz[1],
- hz[end],
- )
- Iv_wy_bc_bf = (; Iv_wy_bc_bf..., B3D = I3D * (Iv_wy_bc_bf.B1D ⊗ I(Ny + 1) ⊗ I(Nx)))
- Iv_wy_bc_bf = (; Iv_wy_bc_bf..., Bbc = I3D * (Iv_wy_bc_bf.Btemp ⊗ I(Ny + 1) ⊗ I(Nx)))
-
- # Boundary conditions low/up
- Nb = Nwy_in + 1 - Nvy_in
- Iv_wy_bc_lu = bc_general(
- Nwy_in + 1,
- Nvy_in,
- Nb,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Iv_wy_bc_lu = (; Iv_wy_bc_lu..., B3D = I(Nz) ⊗ Iv_wy_bc_lu.B1D ⊗ I(Nx))
- Iv_wy_bc_lu = (; Iv_wy_bc_lu..., Bbc = I(Nz) ⊗ Iv_wy_bc_lu.Btemp ⊗ I(Nx))
-
- # Resulting operator:
- Iv_wy = Iv_wy_bc_bf.B3D * Iv_wy_bc_lu.B3D
-
- ## Iw_wz
- diag1 = fill(weight, Nwz_t - 1)
- I1D = spdiagm(Nwz_t - 1, Nwz_t, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Iw_wz_bc = bc_general(
- Nwz_t,
- Nwz_in,
- Nwz_b,
- boundary_conditions.w.z[1],
- boundary_conditions.w.z[2],
- hz[1],
- hz[end],
- )
-
- # Extend to 3D
- Iw_wz = (I1D * Iw_wz_bc.B1D) ⊗ mat_hy ⊗ mat_hx
- Iw_wz_bc = (; Iw_wz_bc..., Bbc = (I1D * Iw_wz_bc.Btemp) ⊗ mat_hy ⊗ mat_hx)
-
- ## Group operators
- (;
- Iu_ux,
- Iv_uy,
- Iw_uz,
- Iu_vx,
- Iv_vy,
- Iw_vz,
- Iu_wx,
- Iv_wy,
- Iw_wz,
- Iu_ux_bc,
- Iv_vy_bc,
- Iw_wz_bc,
- Iv_uy_bc_lr,
- Iv_uy_bc_lu,
- Iu_vx_bc_lr,
- Iu_vx_bc_lu,
- Iw_uz_bc_lr,
- Iw_uz_bc_bf,
- Iw_vz_bc_lu,
- Iw_vz_bc_bf,
- Iu_wx_bc_lr,
- Iu_wx_bc_bf,
- Iv_wy_bc_lu,
- Iv_wy_bc_bf,
- )
-end
diff --git a/src/operators/operator_postprocessing.jl b/src/operators/operator_postprocessing.jl
deleted file mode 100644
index 8ee8869af..000000000
--- a/src/operators/operator_postprocessing.jl
+++ /dev/null
@@ -1,173 +0,0 @@
-"""
- operator_postprocessing(grid, boundary_conditions)
-
-Construct postprocessing operators such as vorticity.
-"""
-operator_postprocessing(grid, boundary_conditions) =
- operator_postprocessing(grid.dimension, grid, boundary_conditions)
-
-# 2D version
-function operator_postprocessing(::Dimension{2}, grid, boundary_conditions)
- (; Nx, Ny, gx, gy, gxd, gyd) = grid
-
- T = eltype(gx)
-
- if all(==(:periodic), (boundary_conditions.u.x[1], boundary_conditions.v.y[1]))
- # For entirely periodic BC, covering entire mesh
-
- # dv/dx, like Sv_vx
- diag = 1 ./ gxd
- ∂x = spdiagm(
- Nx + 1,
- Nx,
- -Nx => diag[[1]],
- -1 => -diag[1:(end-1)],
- 0 => diag[1:(end-1)],
- Nx - 1 => -diag[[end - 1]],
- )
- weight = T(1)
- repeat_x = spdiagm(Nx + 1, Nx, -Nx => [weight], 0 => fill(weight, Nx))
-
- # du/dy, like Su_uy
- diag = 1 ./ gyd
- ∂y = spdiagm(
- Ny + 1,
- Ny,
- -Ny => diag[[1]],
- -1 => -diag[1:(end-1)],
- 0 => diag[1:(end-1)],
- Ny - 1 => -diag[[end - 1]],
- )
- weight = T(1)
- repeat_y = spdiagm(Ny + 1, Ny, -Ny => [weight], 0 => fill(weight, Ny))
- else
- # dv/dx, like Sv_vx
- diag = 1 ./ gy[2:(end-1)]
- ∂y = spdiagm(Ny - 1, Ny, 0 => -diag, 1 => diag)
- repeat_y = I(Ny - 1)
-
- # du/dy, like Su_uy
- diag = 1 ./ gx[2:(end-1)]
- ∂x = spdiagm(Nx - 1, Nx, 0 => -diag, 1 => diag)
- repeat_x = I(Nx - 1)
- end
-
- # Extend to 2D
- Wu_uy = ∂y ⊗ repeat_x
- Wv_vx = repeat_y ⊗ ∂x
-
- (; Wv_vx, Wu_uy)
-end
-
-# 3D version
-function operator_postprocessing(::Dimension{3}, grid, boundary_conditions)
- (; Nx, Ny, Nz, gx, gy, gz, gxd, gyd, gzd) = grid
-
- T = eltype(gx)
-
- if all(
- ==(:periodic),
- (
- boundary_conditions.u.x[1],
- boundary_conditions.v.y[1],
- boundary_conditions.w.z[1],
- ),
- )
- # For entirely periodic BC, covering entire mesh
-
- diag = 1 ./ gxd
- ∂x = spdiagm(
- Nx + 1,
- Nx,
- -Nx => diag[[1]],
- -1 => -diag[1:(end-1)],
- 0 => diag[1:(end-1)],
- Nx - 1 => -diag[[end - 1]],
- )
- # FIXME: nonuniform weights: 1/gi / (1/gi + 1/gj) ?
- weight = T(1 // 2)
- average_x = spdiagm(
- Nx + 1,
- Nx,
- -Nx => [weight],
- -1 => fill(weight, Nx),
- 0 => fill(weight, Nx),
- Nx - 1 => [weight],
- )
- weight = T(1)
- repeat_x = spdiagm(Nx + 1, Nx, -Nx => [weight], 0 => fill(weight, Nx))
-
- diag = 1 ./ gyd
- ∂y = spdiagm(
- Ny + 1,
- Ny,
- -Ny => diag[[1]],
- -1 => -diag[1:(end-1)],
- 0 => diag[1:(end-1)],
- Ny - 1 => -diag[[end - 1]],
- )
- weight = T(1 // 2)
- average_y = spdiagm(
- Ny + 1,
- Ny,
- -Ny => [weight],
- -1 => fill(weight, Ny),
- 0 => fill(weight, Ny),
- Ny - 1 => [weight],
- )
- weight = T(1)
- repeat_y = spdiagm(Ny + 1, Ny, -Ny => [weight], 0 => fill(weight, Ny))
-
- diag = 1 ./ gzd
- ∂z = spdiagm(
- Nz + 1,
- Nz,
- -Nz => diag[[1]],
- -1 => -diag[1:(end-1)],
- 0 => diag[1:(end-1)],
- Nz - 1 => -diag[[end - 1]],
- )
- weight = T(1 // 2)
- average_z = spdiagm(
- Nz + 1,
- Nz,
- -Nz => [weight],
- -1 => fill(weight, Nz),
- 0 => fill(weight, Nz),
- Nz - 1 => [weight],
- )
- weight = T(1)
- repeat_z = spdiagm(Nz + 1, Nz, -Nz => [weight], 0 => fill(weight, Nz))
- else
- diag = 1 ./ gx[2:(end-1)]
- ∂x = spdiagm(Nx - 1, Nx, 0 => -diag, 1 => diag)
- diag = fill(T(1 // 2), Nx - 1)
- average_x = spdiagm(Nx - 1, Nx, 0 => diag, 1 => diag)
- # average_x = sparse(I, Nx - 1, Nx)
- repeat_x = I(Nx - 1)
-
- diag = 1 ./ gy[2:(end-1)]
- ∂y = spdiagm(Ny - 1, Ny, 0 => -diag, 1 => diag)
- diag = fill(T(1 // 2), Ny - 1)
- average_y = spdiagm(Ny - 1, Ny, 0 => diag, 1 => diag)
- # average_y = sparse(I, Ny - 1, Ny)
- repeat_y = I(Ny - 1)
-
- diag = 1 ./ gz[2:(end-1)]
- ∂z = spdiagm(Nz - 1, Nz, 0 => -diag, 1 => diag)
- diag = fill(T(1 / 2), Nz - 1)
- average_z = spdiagm(Nz - 1, Nz, 0 => diag, 1 => diag)
- # average_z = sparse(I, Nz - 1, Nz)
- repeat_z = I(Nz - 1)
- end
-
- # Extend to 3D
- Wu_uy = average_z ⊗ ∂y ⊗ repeat_x
- Wu_uz = ∂z ⊗ average_y ⊗ repeat_x
- Wv_vx = average_z ⊗ repeat_y ⊗ ∂x
- Wv_vz = ∂z ⊗ repeat_y ⊗ average_x
- Ww_wy = repeat_z ⊗ ∂y ⊗ average_x
- Ww_wx = repeat_z ⊗ average_y ⊗ ∂x
-
- (; Wu_uy, Wu_uz, Wv_vx, Wv_vz, Ww_wy, Ww_wx)
-end
diff --git a/src/operators/operator_regularization.jl b/src/operators/operator_regularization.jl
deleted file mode 100644
index ba309a553..000000000
--- a/src/operators/operator_regularization.jl
+++ /dev/null
@@ -1,49 +0,0 @@
-"""
- operator_regularization(grid, operators)
-
-Build regularization matrices.
-"""
-operator_regularization(grid, boundary_conditions) =
- operator_regularization(grid.dimension, grid, boundary_conditions)
-
-# 2D version
-function operator_regularization(::Dimension{2}, grid, operators)
- (; Ω, indu, indv) = grid
- (; Dux, Duy, Dvx, Dvy) = operators
- (; Su_ux, Su_uy, Sv_vx, Sv_vy) = operators
-
- Δ = max_size(grid)
- α_reg = Δ^2 / 16
-
- Ωu⁻¹ = 1 ./ Ω[indu]
- Ωv⁻¹ = 1 ./ Ω[indv]
-
- # Diffusive matrices in finite-difference setting, without viscosity
- Diffu_f = Diagonal(Ωu⁻¹) * (Dux * Su_ux + Duy * Su_uy)
- Diffv_f = Diagonal(Ωv⁻¹) * (Dvx * Sv_vx + Dvy * Sv_vy)
-
- (; Diffu_f, Diffv_f, α_reg)
-end
-
-# 3D version
-function operator_regularization(::Dimension{3}, grid, operators)
- (; Ω, indu, indv, indw) = grid
- (; Dux, Duy, Duz, Dvx, Dvy, Dvz, Dwx, Dwy, Dwz) = operators
- (; Su_ux, Su_uy, Su_uz) = operators
- (; Sv_vx, Sv_vy, Sv_vz) = operators
- (; Sw_wx, Sw_wy, Sw_wz) = operators
-
- Δ = max_size(grid)
- α_reg = Δ^2 / 16
-
- Ωu⁻¹ = 1 ./ Ω[indu]
- Ωv⁻¹ = 1 ./ Ω[indv]
- Ωw⁻¹ = 1 ./ Ω[indw]
-
- # Diffusive matrices in finite-difference setting, without viscosity
- Diffu_f = Diagonal(Ωu⁻¹) * (Dux * Su_ux + Duy * Su_uy + Duz * Su_uz)
- Diffv_f = Diagonal(Ωv⁻¹) * (Dvx * Sv_vx + Dvy * Sv_vy + Dvz * Sv_vz)
- Diffw_f = Diagonal(Ωw⁻¹) * (Dwx * Sw_wx + Dwy * Sw_wy + Dwz * Sw_wz)
-
- (; Diffu_f, Diffv_f, Diffw_f, α_reg)
-end
diff --git a/src/operators/operator_turbulent_diffusion.jl b/src/operators/operator_turbulent_diffusion.jl
deleted file mode 100644
index fd8423ccf..000000000
--- a/src/operators/operator_turbulent_diffusion.jl
+++ /dev/null
@@ -1,309 +0,0 @@
-"""
- operator_turbulent_diffusion(grid, boundary_conditions)
-
-Average (turbulent) viscosity to cell faces: from `ν` at `xp`, `yp` to `ν` at `ux`, `uy`,
-`vx`, `vy` locations.
-"""
-operator_turbulent_diffusion(grid, boundary_conditions) =
- operator_turbulent_diffusion(grid.dimension, grid, boundary_conditions)
-
-# 2D version
-function operator_turbulent_diffusion(::Dimension{2}, grid, boundary_conditions)
- (; Npx, Npy) = grid
- (; Nux_in, Nuy_in, Nvx_in, Nvy_in) = grid
- (; hx, hy, gxd, gyd) = grid
- (; Buvy, Bvux, Bkux, Bkvy) = grid
-
- T = eltype(hx)
-
- # Averaging weight:
- weight = T(1 / 2)
-
- ## Nu to ux positions
-
- A1D = I(Npx + 2)
- A1D = Bkux * A1D
-
- # Boundary conditions for ν; mapping from Npx (k) points to Npx+2 points
- Aν_ux_bc = bc_general_stag(
- Npx + 2,
- Npx,
- 2,
- boundary_conditions.ν.x[1],
- boundary_conditions.ν.x[2],
- hx[1],
- hx[end],
- )
- # Then map back to Nux_in+1 points (ux-faces) with Bkux
-
- # Extend to 2D
- Aν_ux = I(Nuy_in) ⊗ (A1D * Aν_ux_bc.B1D)
- Aν_ux_bc = (; Aν_ux_bc..., Bbc = I(Nuy_in) ⊗ (A1D * Aν_ux_bc.Btemp))
-
- ## Nu to uy positions
-
- # Average in x-direction
- diag1 = fill(weight, Npx + 1)
- A1D = spdiagm(Npx + 1, Npx + 2, 0 => diag1, 1 => diag1)
- # Then map to Nux_in points (like Iv_uy) with Bvux
- A1D = Bvux * A1D
-
- # Calculate average in y-direction; no boundary conditions
- diag1 = fill(weight, Npy + 1)
- A1Dy = spdiagm(Npy + 1, Npy + 2, 0 => diag1, 1 => diag1)
- A2Dy = A1Dy ⊗ I(Nux_in)
-
- # Boundary conditions for ν in x-direction;
- # Mapping from Npx (ν) points to Npx+2 points
- Aν_uy_bc_lr = bc_general_stag(
- Npx + 2,
- Npx,
- 2,
- boundary_conditions.ν.x[1],
- boundary_conditions.ν.x[2],
- hx[1],
- hx[end],
- )
-
- # Extend BC to 2D
- A2D = I(Npy + 2) ⊗ (A1D * Aν_uy_bc_lr.B1D)
-
- # Apply BC in y-direction
- Aν_uy_bc_lu = bc_general_stag(
- Npy + 2,
- Npy,
- 2,
- boundary_conditions.ν.y[1],
- boundary_conditions.ν.y[2],
- hy[1],
- hy[end],
- )
-
- A2Dx = A2D * (Aν_uy_bc_lu.B1D ⊗ I(Npx))
-
- Aν_uy = A2Dy * A2Dx
-
- Aν_uy_bc_lr = (; Aν_uy_bc_lr..., B2D = A2Dy * (I(Npy + 2) ⊗ (A1D * Aν_uy_bc_lr.Btemp)))
- Aν_uy_bc_lu = (; Aν_uy_bc_lu..., B2D = A2Dy * A2D * (Aν_uy_bc_lu.Btemp ⊗ I(Npx)))
-
- ## Nu to vx positions
- diag1 = fill(weight, Npy + 1)
- A1D = spdiagm(Npy + 1, Npy + 2, 0 => diag1, 1 => diag1)
-
- # Map to Nvy_in points (like Iu_vx) with Buvy
- A1D = Buvy * A1D
-
- # Calculate average in x-direction; no boundary conditions
- diag1 = fill(weight, Npx + 1)
- A1Dx = spdiagm(Npx + 1, Npx + 2, 0 => diag1, 1 => diag1)
- A2Dx = kron(I(Nvy_in), A1Dx)
-
- # Boundary conditions for ν in y-direction;
- # Mapping from Npy (ν) points to Npy+2 points
- Aν_vx_bc_lu = bc_general_stag(
- Npy + 2,
- Npy,
- 2,
- boundary_conditions.ν.y[1],
- boundary_conditions.ν.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend BC to 2D
- A2D = (A1D * Aν_vx_bc_lu.B1D) ⊗ I(Npx + 2)
-
- # Apply boundary conditions also in x-direction:
- Aν_vx_bc_lr = bc_general_stag(
- Npx + 2,
- Npx,
- 2,
- boundary_conditions.ν.x[1],
- boundary_conditions.ν.x[2],
- hx[1],
- hx[end],
- )
-
- A2Dy = A2D * (I(Npy) ⊗ Aν_vx_bc_lr.B1D)
-
- Aν_vx = A2Dx * A2Dy
-
- Aν_vx_bc_lu = (; Aν_vx_bc_lu..., B2D = A2Dx * ((A1D * Aν_vx_bc_lu.Btemp) ⊗ I(Npx + 2)))
- Aν_vx_bc_lr = (; Aν_vx_bc_lr..., B2D = A2Dx * A2D * (I(Npy) ⊗ Aν_vx_bc_lr.Btemp))
-
- ## Nu to vy positions
- A1D = I(Npy + 2)
- # Then map back to Nvy_in+1 points (vy-faces) with Bkvy
- A1D = Bkvy * A1D
-
- # Boundary conditions for ν; mapping from Npy (ν) points to Npy+2 (vy faces) points
- Aν_vy_bc = bc_general_stag(
- Npy + 2,
- Npy,
- 2,
- boundary_conditions.ν.y[1],
- boundary_conditions.ν.y[2],
- hy[1],
- hy[end],
- )
-
- # Extend to 2D
- Aν_vy = (A1D * Aν_vy_bc.B1D) ⊗ I(Nvx_in)
- Aν_vy_bc = (; Aν_vy_bc..., Bbc = ((A1D * Aν_vy_bc.Btemp) ⊗ I(Nvx_in)))
-
- # So ν at vy is given by:
- # (Aν_vy * k + yAν_vy)
-
- ## Get derivatives u_x, u_y, v_x and v_y at cell centers
- # Differencing velocity to ν-points
-
- ## Du/dx
-
- # Differencing matrix
- diag1 = 1 ./ hx
- C1D = spdiagm(Npx, Npx + 1, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Cux_k_bc = bc_general(
- Npx + 1,
- Nux_in,
- Npx + 1 - Nux_in,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- Cux_k = I(Npy) ⊗ (C1D * Cux_k_bc.B1D)
- Cux_k_bc = (; Cux_k_bc..., Bbc = I(Npy) ⊗ (C1D * Cux_k_bc.Btemp))
-
- # Cux_k*uₕ+yCux_k;
-
- ## Du/dy
-
- # Average to k-positions (in x-dir)
- weight = T(1 / 2)
- diag1 = fill(weight, Npx)
- A1D = spdiagm(Npx, Npx + 1, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Auy_k_bc = bc_general(
- Npx + 1,
- Nux_in,
- Npx + 1 - Nux_in,
- boundary_conditions.u.x[1],
- boundary_conditions.u.x[2],
- hx[1],
- hx[end],
- )
-
- Auy_k = I(Npy) ⊗ (A1D * Auy_k_bc.B1D)
- Auy_k_bc = (; Auy_k_bc..., Bbc = I(Npy) ⊗ (A1D * Auy_k_bc.Btemp))
-
- # Take differences
- gydnew = gyd[1:end-1] + gyd[2:end] # Differencing over 2*Δy
- diag2 = 1 ./ gydnew
- C1D = spdiagm(Npy, Npy + 2, 0 => -diag2, 2 => diag2)
-
- Cuy_k_bc = bc_general_stag(
- Npy + 2,
- Npy,
- 2,
- boundary_conditions.u.y[1],
- boundary_conditions.u.y[2],
- hy[1],
- hy[end],
- )
-
- Cuy_k = (C1D * Cuy_k_bc.B1D) ⊗ I(Npx)
- Cuy_k_bc = (; Cuy_k_bc..., Bbc = (C1D * Cuy_k_bc.Btemp) ⊗ I(Npx))
-
- ## Dv/dx
-
- # Average to k-positions (in y-dir)
- weight = T(1 / 2)
- diag1 = fill(weight, Npy)
- A1D = spdiagm(Npy, Npy + 1, 0 => diag1, 1 => diag1)
-
- # Boundary conditions
- Avx_k_bc = bc_general(
- Npy + 1,
- Nvy_in,
- Npy + 1 - Nvy_in,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
- Avx_k = (A1D * Avx_k_bc.B1D) ⊗ I(Npx)
- Avx_k_bc = (; Avx_k_bc..., Bbc = (A1D * Avx_k_bc.Btemp) ⊗ I(Npx))
-
- # Take differences
- gxdnew = gxd[1:end-1] + gxd[2:end] # Differencing over 2*Δx
- diag2 = 1 ./ gxdnew
- C1D = spdiagm(Npx, Npx + 2, 0 => -diag2, 2 => diag2)
-
- Cvx_k_bc = bc_general_stag(
- Npx + 2,
- Npx,
- Npx + 2 - Npx,
- boundary_conditions.v.x[1],
- boundary_conditions.v.x[2],
- hx[1],
- hx[end],
- )
-
- Cvx_k = I(Npy) ⊗ (C1D * Cvx_k_bc.B1D)
- Cvx_k_bc = (; Cvx_k_bc..., Bbc = I(Npy) ⊗ (C1D * Cvx_k_bc.Btemp))
-
- ## Dv/dy
-
- # Differencing matrix
- diag1 = 1 ./ hy
- C1D = spdiagm(Npy, Npy + 1, 0 => -diag1, 1 => diag1)
-
- # Boundary conditions
- Cvy_k_bc = bc_general(
- Npy + 1,
- Nvy_in,
- Npy + 1 - Nvy_in,
- boundary_conditions.v.y[1],
- boundary_conditions.v.y[2],
- hy[1],
- hy[end],
- )
-
- Cvy_k = (C1D * Cvy_k_bc.B1D) ⊗ I(Npx)
- Cvy_k_bc = (; Cvy_k_bc..., Bbc = (C1D * Cvy_k_bc.Btemp) ⊗ I(Npx))
-
- ## Group operators
- (;
- Aν_ux,
- Aν_ux_bc,
- Aν_uy,
- Aν_uy_bc_lr,
- Aν_uy_bc_lu,
- Aν_vx,
- Aν_vx_bc_lr,
- Aν_vx_bc_lu,
- Aν_vy,
- Aν_vy_bc,
- Cux_k,
- Cux_k_bc,
- Cuy_k,
- Cuy_k_bc,
- Cvx_k,
- Cvx_k_bc,
- Cvy_k,
- Cvy_k_bc,
- Auy_k,
- Auy_k_bc,
- Avx_k,
- Avx_k_bc,
- )
-end
-
-# TODO: 3D implementation
-function operator_turbulent_diffusion(::Dimension{3}, grid, boundary_conditions)
- error("Not implemented (3D)")
-end
diff --git a/src/operators/operator_viscosity.jl b/src/operators/operator_viscosity.jl
deleted file mode 100644
index d066c54db..000000000
--- a/src/operators/operator_viscosity.jl
+++ /dev/null
@@ -1,16 +0,0 @@
-"""
- operator_viscosity(viscosity_model, grid, boundary_conditions)
-
-Classical turbulence modelling via the diffusive term
-"""
-function operator_viscosity end
-
-operator_viscosity(::LaminarModel, grid, boundary_conditions) = (;)
-
-function operator_viscosity(
- ::Union{QRModel,SmagorinskyModel,MixingLengthModel},
- grid,
- boundary_conditions,
-)
- operator_turbulent_diffusion(grid, boundary_conditions)
-end
diff --git a/src/operators/operators.jl b/src/operators/operators.jl
deleted file mode 100644
index 97a7610be..000000000
--- a/src/operators/operators.jl
+++ /dev/null
@@ -1,29 +0,0 @@
-"""
- Operators(grid, boundary_conditions, viscosity_model)
-
-Build operators.
-"""
-function Operators(grid, boundary_conditions, viscosity_model)
- # Averaging operators
- op_ave = operator_averaging(grid, boundary_conditions)
-
- # Interpolation operators
- op_int = operator_interpolation(grid, boundary_conditions)
-
- # Divergence (u, v) -> p and gradient p -> (u, v) operator
- op_div = operator_divergence(grid, boundary_conditions)
-
- # Convection operators on u- and v- centered volumes
- op_con = operator_convection_diffusion(grid, boundary_conditions)
-
- # Regularization modelling - this changes the convective term
- op_reg = operator_regularization(grid, op_con)
-
- # Classical turbulence modelling via the diffusive term
- op_vis = operator_viscosity(viscosity_model, grid, boundary_conditions)
-
- # Post-processing
- op_pos = operator_postprocessing(grid, boundary_conditions)
-
- (; op_ave..., op_int..., op_div..., op_con..., op_reg..., op_vis..., op_pos...)
-end
diff --git a/src/postprocess/get_streamfunction.jl b/src/postprocess/get_streamfunction.jl
deleted file mode 100644
index 434f81c71..000000000
--- a/src/postprocess/get_streamfunction.jl
+++ /dev/null
@@ -1,110 +0,0 @@
-"""
- get_streamfunction(setup, V, t)
-
-Compute streamfunction ``\\psi`` from a Poisson equation ``\\nabla^2 \\psi = -\\omega``.
-"""
-function get_streamfunction end
-
-get_streamfunction(setup, V, t) = get_streamfunction(setup.grid.dimension, setup, V, t)
-
-# 2D version
-function get_streamfunction(::Dimension{2}, setup, V, t)
- (; grid, operators, boundary_conditions) = setup
- (; u_bc, v_bc) = boundary_conditions
- (; indu, indv, Nux_in, Nvx_in, Nx, Ny) = grid
- (; hx, hy, x, y, xp, yp) = grid
- (; Wv_vx, Wu_uy) = operators
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
-
- # Boundary values by integrating around domain
- # Start with ψ = 0 at lower left corner
-
- # u = d ψ / dy; integrate low->up
- if boundary_conditions.u.x[1] == :dirichlet
- # u1 = interp1(y, uLe, yp);
- u1 = u_bc.(x[1], yp, t)
- elseif boundary_conditions.u.x[1] ∈ (:pressure, :periodic)
- u1 = uₕ[1:Nux_in:end]
- end
- ψLe = cumsum(hy .* u1)
- ψUpLe = ψLe[end]
- ψLe = ψLe[1:(end-1)]
-
- # v = -d ψ / dx; integrate left->right
- if boundary_conditions.v.y[2] == :dirichlet
- v1 = v_bc.(xp, y[end], t)
- elseif boundary_conditions.v.y[2] == :pressure
- v1 = vₕ[(end-Nvx_in+1):end]
- elseif boundary_conditions.v.y[2] == :periodic
- v1 = vₕ[1:Nvx_in]
- end
- ψUp = ψUpLe .- cumsum(hx .* v1)
- ψUpRi = ψUp[end]
- ψUp = ψUp[1:(end-1)]
-
- # u = d ψ / dy; integrate up->lo
- if boundary_conditions.u.x[2] == :dirichlet
- u2 = u_bc.(x[end], yp, t)
- elseif boundary_conditions.u.x[2] == :pressure
- u2 = uₕ[Nux_in:Nux_in:end]
- elseif boundary_conditions.u.x[2] == :periodic
- u2 = uₕ[1:Nux_in:end]
- end
- ψRi = ψUpRi .- cumsum(hy[end:-1:1] .* u2[end:-1:1])
- ψLoRi = ψRi[end]
- ψRi = ψRi[(end-1):-1:1]
-
- # v = -d ψ / dx; integrate right->left
- if boundary_conditions.v.y[1] == :dirichlet
- v2 = v_bc.(xp, y[1], t)
- elseif boundary_conditions.v.y[1] ∈ (:pressure, :periodic)
- v2 = vₕ[1:Nvx_in]
- end
- ψLo = ψLoRi .+ cumsum(hx[end:-1:1] .* v2[end:-1:1])
- ψLoLe = ψLo[end]
- ψLo = ψLo[(end-1):-1:1]
-
- abs(ψLoLe) > 1e-12 && @warn "Contour integration of ψ not consistent" abs(ψLoLe)
-
- # Solve del^2 ψ = -ω
- # Only dirichlet boundary conditions because we calculate streamfunction at
- # Inner points only
-
- # X-direction
- diag = 1 ./ hx
- Q1D = spdiagm(Nx, Nx + 1, 0 => -diag, 1 => diag)
- Qx_bc = bc_general(Nx + 1, Nx - 1, 2, :dirichlet, :dirichlet, hx[1], hx[end])
-
- # Extend to 2D
- Q2Dx = I(Ny - 1) ⊗ (Q1D * Qx_bc.B1D)
- yQx = (I(Ny - 1) ⊗ (Q1D * Qx_bc.Btemp)) * (ψLe ⊗ Qx_bc.ybc1 + ψRi ⊗ Qx_bc.ybc2)
-
- # Y-direction
- diag = 1 ./ hy
- Q1D = spdiagm(Ny, Ny + 1, 0 => -diag, 1 => diag)
- Qy_bc = bc_general(Ny + 1, Ny - 1, 2, :dirichlet, :dirichlet, hy[1], hy[end])
-
- # Extend to 2D
- Q2Dy = (Q1D * Qy_bc.B1D) ⊗ I(Nx - 1)
- yQy = ((Q1D * Qy_bc.Btemp) ⊗ I(Nx - 1)) * (Qy_bc.ybc1 ⊗ ψLo + Qy_bc.ybc2 ⊗ ψUp)
-
- # FIXME: Dimension error in periodic case
- # @show size(Wv_vx) size(Q2Dx) size(Wu_uy) size(Q2Dy)
- Aψ = Wv_vx * Q2Dx + Wu_uy * Q2Dy
- yAψ = Wv_vx * yQx + Wu_uy * yQy
-
- ω = get_vorticity(setup, V, t)
- ω_flat = reshape(ω, length(ω))
-
- # Solve streamfunction from Poisson equation
- ψ = -Aψ \ (ω_flat + yAψ)
-
- reshape(ψ, size(ω)...)
-end
-
-# 3D version
-function get_streamfunction(::Dimension{3}, setup, V, t)
- error("Not implemented")
-end
diff --git a/src/postprocess/get_velocity.jl b/src/postprocess/get_velocity.jl
deleted file mode 100644
index 094ef5961..000000000
--- a/src/postprocess/get_velocity.jl
+++ /dev/null
@@ -1,49 +0,0 @@
-"""
- get_velocity(setup, V, t)
-
-Get velocity values at pressure points. Interpolate velocities to pressure positions using
-`BMx` and `BMy` (and `BMz`), constructed in operator_divergence.jl.
-"""
-function get_velocity end
-
-get_velocity(setup, V, t) = get_velocity(setup.grid.dimension, setup, V, t)
-
-# 2D version
-function get_velocity(::Dimension{2}, setup, V, t)
- (; grid, operators) = setup
- (; Npx, Npy, indu, indv) = grid
- (; Au_ux, Av_vy, Bup, Bvp) = operators
-
- # Evaluate boundary conditions at current time
- bc_vectors = get_bc_vectors(setup, t)
- (; yAu_ux, yAv_vy) = bc_vectors
-
- uh = @view V[indu]
- vh = @view V[indv]
-
- up = reshape(Bup * (Au_ux * uh + yAu_ux), Npx, Npy)
- vp = reshape(Bvp * (Av_vy * vh + yAv_vy), Npx, Npy)
-
- up, vp
-end
-
-# 3D version
-function get_velocity(::Dimension{3}, setup, V, t)
- (; grid, operators) = setup
- (; Au_ux, Av_vy, Aw_wz, Bup, Bvp, Bwp) = operators
- (; Npx, Npy, Npz, indu, indv, indw) = grid
-
- # Evaluate boundary conditions at current time
- bc_vectors = get_bc_vectors(setup, t)
- (; yAu_ux, yAv_vy, yAw_wz) = bc_vectors
-
- uh = @view V[indu]
- vh = @view V[indv]
- wh = @view V[indw]
-
- up = reshape(Bup * (Au_ux * uh + yAu_ux), Npx, Npy, Npz)
- vp = reshape(Bvp * (Av_vy * vh + yAv_vy), Npx, Npy, Npz)
- wp = reshape(Bwp * (Aw_wz * wh + yAw_wz), Npx, Npy, Npz)
-
- up, vp, wp
-end
diff --git a/src/postprocess/get_vorticity.jl b/src/postprocess/get_vorticity.jl
deleted file mode 100644
index 72ff1d195..000000000
--- a/src/postprocess/get_vorticity.jl
+++ /dev/null
@@ -1,156 +0,0 @@
-"""
- get_vorticity(setup, V, t)
-
-Get vorticity from velocity field.
-"""
-function get_vorticity end
-
-get_vorticity(setup, V, t) = get_vorticity(setup.grid.dimension, setup, V, t)
-
-# 2D version
-function get_vorticity(::Dimension{2}, setup, V, t)
- (; grid, boundary_conditions) = setup
- (; Nx, Ny) = grid
-
- if all(==(:periodic), (boundary_conditions.u.x[1], boundary_conditions.v.y[1]))
- Nωx = Nx + 1
- Nωy = Ny + 1
- else
- Nωx = Nx - 1
- Nωy = Ny - 1
- end
-
- ω = similar(V, Nωx, Nωy)
-
- vorticity!(ω, setup, V, t)
-end
-
-# 3D version
-function get_vorticity(::Dimension{3}, setup, V, t)
- (; grid, boundary_conditions) = setup
- (; Nx, Ny, Nz) = grid
-
- if all(
- ==(:periodic),
- (
- boundary_conditions.u.x[1],
- boundary_conditions.v.y[1],
- boundary_conditions.w.z[1],
- ),
- )
- Nωx = Nx + 1
- Nωy = Ny + 1
- Nωz = Nz + 1
- else
- Nωx = Nx - 1
- Nωy = Ny - 1
- Nωz = Nz - 1
- end
-
- ω = similar(V, Nωx, Nωy, Nωz)
-
- vorticity!(ω, setup, V, t)
-end
-
-"""
- vorticity!(ω, V, t, setup)
-
-Compute vorticity values at pressure midpoints.
-This should be consistent with `operator_postprocessing.jl`.
-"""
-function vorticity! end
-
-vorticity!(ω, setup, V, t) = vorticity!(setup.grid.dimension, ω, setup, V, t)
-
-# 2D version
-function vorticity!(::Dimension{2}, ω, setup, V, t)
- (; grid, operators, boundary_conditions) = setup
- (; indu, indv, Nux_in, Nvy_in, Nx, Ny) = grid
- (; Wv_vx, Wu_uy) = operators
-
- T = eltype(V)
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- ω_flat = reshape(ω, :)
-
- if boundary_conditions.u.x[1] == :periodic && boundary_conditions.v.y[1] == :periodic
- uₕ_in = uₕ
- vₕ_in = vₕ
- else
- # Velocity at inner points
- diagpos = 0
- boundary_conditions.u.x[1] == :pressure && (diagpos = 1)
- boundary_conditions.u.x[1] == :periodic && (diagpos = 1)
- B1D = spdiagm(Nx - 1, Nux_in, diagpos => ones(T, Nx - 1))
- B2D = I(Ny) ⊗ B1D
- uₕ_in = B2D * uₕ
-
- diagpos = 0
- boundary_conditions.v.y[1] == :pressure && (diagpos = 1)
- boundary_conditions.v.y[1] == :periodic && (diagpos = 1)
- B1D = spdiagm(Ny - 1, Nvy_in, diagpos => ones(T, Ny - 1))
- B2D = B1D ⊗ I(Nx)
- vₕ_in = B2D * vₕ
- end
-
- # ω_flat .= Wv_vx * vₕ_in - Wu_uy * uₕ_in
- mul!(ω_flat, Wv_vx, vₕ_in) # a = b * c
- mul!(ω_flat, Wu_uy, uₕ_in, -1, 1) # a = -b * c + a
-
- ω
-end
-
-# 3D version
-function vorticity!(::Dimension{3}, ω, setup, V, t)
- (; grid, operators, boundary_conditions) = setup
- (; indu, indv, indw, Nux_in, Nvy_in, Nwz_in, Nx, Ny, Nz) = grid
- (; Wu_uy, Wu_uz, Wv_vx, Wv_vz, Ww_wx, Ww_wy) = operators
-
- T = eltype(V)
-
- uₕ = @view V[indu]
- vₕ = @view V[indv]
- wₕ = @view V[indw]
- ω_flat = reshape(ω, :)
-
- if boundary_conditions.u.x[1] == :periodic &&
- boundary_conditions.v.y[1] == :periodic &&
- boundary_conditions.w.z[1] == :periodic
- uₕ_in = uₕ
- vₕ_in = vₕ
- wₕ_in = wₕ
- else
- # Velocity at inner points
- diagpos = 0
- boundary_conditions.u.x[1] == :pressure && (diagpos = 1)
- boundary_conditions.u.x == (:periodic, :periodic) && (diagpos = 1)
- B1D = spdiagm(Nx - 1, Nux_in, diagpos => ones(T, Nx - 1))
- B2D = I(Nz) ⊗ I(Ny) ⊗ B1D
- uₕ_in = B2D * uₕ
-
- diagpos = 0
- boundary_conditions.v.y[1] == :pressure && (diagpos = 1)
- boundary_conditions.v.y == (:periodic, :periodic) && (diagpos = 1)
- B1D = spdiagm(Ny - 1, Nvy_in, diagpos => ones(T, Ny - 1))
- B2D = I(Nz) ⊗ B1D ⊗ I(Nx)
- vₕ_in = B2D * vₕ
-
- diagpos = 0
- boundary_conditions.w.z[1] == :pressure && (diagpos = 1)
- boundary_conditions.w.z == (:periodic, :periodic) && (diagpos = 1)
- B1D = spdiagm(Nz - 1, Nwz_in, diagpos => ones(T, Nz - 1))
- B2D = B1D ⊗ I(Ny) ⊗ I(Nx)
- wₕ_in = B2D * wₕ
- end
-
- # ωx_flat .= Ww_wy * wₕ_in - Wy_yz * vₕ_in
- # ωy_flat .= Wu_uz * uₕ_in - Ww_wx * wₕ_in
- # ωz_flat .= Wv_vx * vₕ_in - Wu_uy * uₕ_in
-
- ω_flat .= .√(
- (Ww_wy * wₕ_in .- Wv_vz * vₕ_in) .^ 2 .+ (Wu_uz * uₕ_in .- Ww_wx * wₕ_in) .^ 2 .+ (Wv_vx * vₕ_in .- Wu_uy * uₕ_in) .^ 2,
- )
-
- ω
-end
diff --git a/src/postprocess/plot_force.jl b/src/postprocess/plot_force.jl
deleted file mode 100644
index 17fa505c1..000000000
--- a/src/postprocess/plot_force.jl
+++ /dev/null
@@ -1,53 +0,0 @@
-"""
- plot_force(setup, t; kwargs...)
-
-Plot body force.
-"""
-function plot_force end
-
-plot_force(setup, t; kwargs...) = plot_force(setup.grid.dimension, setup, t; kwargs...)
-
-# 2D version
-function plot_force(::Dimension{2}, setup, t; kwargs...)
- (; grid, force) = setup
- (; xp, yp, xlims, ylims) = grid
- (; xu, yu, indu, xlims, ylims) = grid
- F = force
-
- Fp = reshape(F[indu], size(xu))
- # TODO: norm of F instead of Fu
- xp, yp = xu[:, 1], yu[1, :]
-
- # Levels
- μ, σ = mean(Fp), std(Fp)
- ≈(μ + σ, μ; rtol = 1e-8, atol = 1e-8) && (σ = 1e-4)
- levels = LinRange(μ - 1.5σ, μ + 1.5σ, 10)
-
- fig = Figure()
- ax = Axis(fig[1, 1]; aspect = DataAspect(), title = "Force", xlabel = "x", ylabel = "y")
- limits!(ax, xlims[1], xlims[2], ylims[1], ylims[2])
- cf = contourf!(ax, xp, yp, Fp; extendlow = :auto, extendhigh = :auto, levels, kwargs...)
- Colorbar(fig[1, 2], cf)
- fig
-end
-
-# 3D version
-function plot_force(::Dimension{3}, setup, t; kwargs...)
- (; grid, force) = setup
- (; xu, yu, zu, indu) = grid
- F = force
-
- # Get force at pressure points
- # up, vp, wp = get_velocity(setup, V, t)
- # qp = map((u, v, w) -> √sum(u^2 + v^2 + w^2), up, vp, wp)
- Fp = reshape(F[indu], size(xu))
- # TODO: norm of F instead of Fu
- xp, yp, zp = xu[:, 1, 1], yu[1, :, 1], zu[1, 1, :]
-
- # Levels
- μ, σ = mean(Fp), std(Fp)
- levels = LinRange(μ - 3σ, μ + 3σ, 10)
-
- # contour(xp, yp, zp, Fp; levels, kwargs...)
- contour(xp, yp, zp, Fp)
-end
diff --git a/src/postprocess/plot_pressure.jl b/src/postprocess/plot_pressure.jl
deleted file mode 100644
index 1a4bbbb56..000000000
--- a/src/postprocess/plot_pressure.jl
+++ /dev/null
@@ -1,56 +0,0 @@
-"""
- plot_pressure(setup, p; kwargs...)
-
-Plot pressure.
-"""
-function plot_pressure end
-
-plot_pressure(setup, p; kwargs...) =
- plot_pressure(setup.grid.dimension, setup, p; kwargs...)
-
-# 2D version
-function plot_pressure(::Dimension{2}, setup, p; kwargs...)
- (; Nx, Ny, Npx, Npy, xp, yp, xlims, ylims) = setup.grid
-
- # Reshape
- p = reshape(p, Npx, Npy)
-
- # Levels
- μ, σ = mean(p), std(p)
- ≈(μ + σ, μ; rtol = 1e-8, atol = 1e-8) && (σ = 1e-4)
- levels = LinRange(μ - 1.5σ, μ + 1.5σ, 10)
-
- # Plot pressure
- fig = Figure()
- ax = Axis(
- fig[1, 1];
- aspect = DataAspect(),
- title = "Pressure",
- xlabel = "x",
- ylabel = "y",
- )
- limits!(ax, xlims[1], xlims[2], ylims[1], ylims[2])
- cf = contourf!(ax, xp, yp, p; extendlow = :auto, extendhigh = :auto, levels, kwargs...)
- # Colorbar(fig[1,2], cf)
- Colorbar(fig[1, 2], cf)
-
- # save("output/pressure.png", fig, pt_per_unit = 2)
-
- fig
-end
-
-# 3D version
-function plot_pressure(::Dimension{3}, setup, p; kwargs...)
- (; Nx, Ny, Nz, Npx, Npy, Npz, xp, yp, zp) = setup.grid
-
- # Reshape
- p = reshape(p, Npx, Npy, Npz)
-
- # Levels
- μ, σ = mean(p), std(p)
- levels = LinRange(μ - 5σ, μ + 5σ, 10)
-
- contour(xp, yp, zp, p; levels, kwargs...)
-
- # save("output/pressure.png", fig, pt_per_unit = 2)
-end
diff --git a/src/postprocess/plot_streamfunction.jl b/src/postprocess/plot_streamfunction.jl
deleted file mode 100644
index 8bb1dd2f5..000000000
--- a/src/postprocess/plot_streamfunction.jl
+++ /dev/null
@@ -1,61 +0,0 @@
-"""
- plot_streamfunction(setup, V, t; kwargs...)
-
-Plot streamfunction.
-"""
-function plot_streamfunction end
-
-plot_streamfunction(setup, V, t; kwargs...) =
- plot_streamfunction(setup.grid.dimension, setup, V, t; kwargs...)
-
-# 2D version
-function plot_streamfunction(::Dimension{2}, setup, V, t; kwargs...)
- (; grid, boundary_conditions) = setup
- (; x, y, xlims, ylims) = grid
-
- if all(==(:periodic), (boundary_conditions.u.x[1], boundary_conditions.v.y[1]))
- xψ = x
- yψ = y
- else
- xψ = x[2:(end-1)]
- yψ = y[2:(end-1)]
- end
-
- # Get fields
- ψ = get_streamfunction(setup, V, t)
-
- # Levels
- μ, σ = mean(ψ), std(ψ)
- ≈(μ + σ, μ; rtol = 1e-8, atol = 1e-8) && (σ = 1e-4)
- levels = LinRange(μ - 1.5σ, μ + 1.5σ, 10)
-
- # Plot stream function
- fig = Figure()
- ax = Axis(
- fig[1, 1],
- aspect = DataAspect(),
- title = "Stream function",
- xlabel = "x",
- ylabel = "y",
- )
- limits!(ax, xlims[1], xlims[2], ylims[1], ylims[2])
- cf = contourf!(
- ax,
- xψ,
- yψ,
- ψ;
- extendlow = :auto,
- extendhigh = :auto,
- # levels,
- kwargs...,
- )
- Colorbar(fig[1, 2], cf)
- # save("output/streamfunction.png", fig, pt_per_unit = 2)
-
- fig
-end
-
-# 3D version
-function plot_streamfunction(::Dimension{3}, setup, V, t; kwargs...)
- error("Not implemented (3D)")
-end
diff --git a/src/postprocess/plot_velocity.jl b/src/postprocess/plot_velocity.jl
deleted file mode 100644
index e37deee90..000000000
--- a/src/postprocess/plot_velocity.jl
+++ /dev/null
@@ -1,55 +0,0 @@
-"""
- plot_velocity(setup, V, t; kwargs...)
-
-Plot velocity.
-"""
-function plot_velocity end
-
-plot_velocity(setup, V, t; kwargs...) =
- plot_velocity(setup.grid.dimension, setup, V, t; kwargs...)
-
-# 2D version
-function plot_velocity(::Dimension{2}, setup, V, t; kwargs...)
- (; xp, yp, xlims, ylims) = setup.grid
-
- # Get velocity at pressure points
- up, vp = get_velocity(setup, V, t)
- qp = map((u, v) -> √(u^2 + v^2), up, vp)
-
- # Levels
- μ, σ = mean(qp), std(qp)
- ≈(μ + σ, μ; rtol = 1e-8, atol = 1e-8) && (σ = 1e-4)
- levels = LinRange(μ - 1.5σ, μ + 1.5σ, 10)
-
- fig = Figure()
- ax = Axis(
- fig[1, 1];
- aspect = DataAspect(),
- title = "Velocity",
- xlabel = "x",
- ylabel = "y",
- )
- limits!(ax, xlims[1], xlims[2], ylims[1], ylims[2])
- cf = contourf!(ax, xp, yp, qp; extendlow = :auto, extendhigh = :auto, levels, kwargs...)
- Colorbar(fig[1, 2], cf)
- # Colorbar(fig[2,1], cf; vertical = false)
- fig
-end
-
-# 3D version
-function plot_velocity(::Dimension{3}, setup, V, t; kwargs...)
- (; xu, yu, zu, indu) = setup.grid
-
- # Get velocity at pressure points
- # up, vp, wp = get_velocity(setup, V, t)
- # qp = map((u, v, w) -> √sum(u^2 + v^2 + w^2), up, vp, wp)
- qp = reshape(V[indu], size(xu))
- xp, yp, zp = xu[:, 1, 1], yu[1, :, 1], zu[1, 1, :]
-
- # Levels
- μ, σ = mean(qp), std(qp)
- levels = LinRange(μ - 3σ, μ + 3σ, 10)
-
- # contour(xp, yp, zp, qp; levels, kwargs...)
- contour(xp, yp, zp, qp; kwargs...)
-end
diff --git a/src/postprocess/plot_vorticity.jl b/src/postprocess/plot_vorticity.jl
deleted file mode 100644
index 1d8a50e76..000000000
--- a/src/postprocess/plot_vorticity.jl
+++ /dev/null
@@ -1,79 +0,0 @@
-"""
- plot_vorticity(setup, V, t; kwargs...)
-
-Plot vorticity field.
-"""
-function plot_vorticity end
-
-plot_vorticity(setup, V, t; kwargs...) =
- plot_vorticity(setup.grid.dimension, setup, V, t; kwargs...)
-
-# 2D version
-function plot_vorticity(::Dimension{2}, setup, V, t; kwargs...)
- (; grid, boundary_conditions) = setup
- (; x, y, xlims, ylims) = grid
-
- if all(==(:periodic), (boundary_conditions.u.x[1], boundary_conditions.v.y[1]))
- xω = x
- yω = y
- else
- xω = x[2:(end-1)]
- yω = y[2:(end-1)]
- end
-
- # Get fields
- ω = get_vorticity(setup, V, t)
-
- # Levels
- μ, σ = mean(ω), std(ω)
- ≈(μ + σ, μ; rtol = 1e-8, atol = 1e-8) && (σ = 1e-4)
- levels = LinRange(μ - 1.5σ, μ + 1.5σ, 10)
-
- # Plot vorticity
- fig = Figure()
- ax = Axis(
- fig[1, 1];
- aspect = DataAspect(),
- title = "Vorticity",
- xlabel = "x",
- ylabel = "y",
- )
- limits!(ax, xlims[1], xlims[2], ylims[1], ylims[2])
- cf = contourf!(ax, xω, yω, ω; extendlow = :auto, extendhigh = :auto, levels, kwargs...)
- Colorbar(fig[1, 2], cf)
-
- # save("output/vorticity.png", fig, pt_per_unit = 2)
-
- fig
-end
-
-# 3D version
-function plot_vorticity(::Dimension{3}, setup, V, t; kwargs...)
- (; grid, boundary_conditions) = setup
- (; x, y, z) = grid
-
- if all(
- ==(:periodic),
- (
- boundary_conditions.u.x[1],
- boundary_conditions.v.y[1],
- boundary_conditions.w.z[1],
- ),
- )
- xω = x
- yω = y
- zω = z
- else
- xω = x[2:(end-1)]
- yω = y[2:(end-1)]
- zω = z[2:(end-1)]
- end
-
- ω = get_vorticity(setup, V, t)
-
- # Levels
- μ, σ = mean(ω), std(ω)
- levels = LinRange(μ - 3σ, μ + 3σ, 10)
-
- contour(xω, yω, zω, ω; levels, kwargs...)
-end
diff --git a/src/postprocess/save_vtk.jl b/src/postprocess/save_vtk.jl
deleted file mode 100644
index de11ebee5..000000000
--- a/src/postprocess/save_vtk.jl
+++ /dev/null
@@ -1,32 +0,0 @@
-"""
- save_vtk(setup, V, p, t, filename = "output/solution")
-
-Save velocity and pressure field to a VTK file.
-
-In the case of a 2D setup, the velocity field is saved as a 3D vector with a
-z-component of zero, as this seems to be preferred by ParaView.
-"""
-function save_vtk(setup, V, p, t, filename = "output/solution")
- parts = split(filename, "/")
- path = join(parts[1:end-1], "/")
- isdir(path) || mkpath(path)
- (; grid) = setup
- N = setup.grid.dimension()
- if N == 2
- (; xp, yp) = grid
- coords = (xp, yp)
- elseif N == 3
- (; xp, yp, zp) = grid
- coords = (xp, yp, zp)
- end
- vtk_grid(filename, coords...) do vtk
- vels = get_velocity(setup, V, t)
- if N == 2
- # ParaView prefers 3D vectors. Add zero z-component.
- wp = zero(vels[1])
- vels = (vels..., wp)
- end
- vtk["velocity"] = vels
- vtk["pressure"] = p
- end
-end
diff --git a/src/preprocess/create_initial_conditions.jl b/src/preprocess/create_initial_conditions.jl
deleted file mode 100644
index 3cd132a6c..000000000
--- a/src/preprocess/create_initial_conditions.jl
+++ /dev/null
@@ -1,253 +0,0 @@
-"""
- create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- [initial_velocity_w,]
- t;
- initial_pressure = nothing,
- pressure_solver = DirectPressureSolver(setup),
- )
-
-Create initial vectors at starting time `t`. If `p_initial` is a function instead of
-`nothing`, calculate compatible IC for the pressure.
-"""
-function create_initial_conditions end
-
-# 2D version
-function create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- t;
- initial_pressure = nothing,
- pressure_solver = DirectPressureSolver(setup),
-)
- (; grid, operators) = setup
- (; xu, yu, xv, yv, xpp, ypp, Ω) = grid
- (; G, M) = operators
-
- T = eltype(xu)
-
- # Boundary conditions
- bc_vectors = get_bc_vectors(setup, t)
- (; yM) = bc_vectors
-
- # Allocate velocity and pressure
- u = zero(xu)[:]
- v = zero(xv)[:]
- p = zero(xpp)[:]
-
- # Initial velocities
- u .= initial_velocity_u.(xu, yu)[:]
- v .= initial_velocity_v.(xv, yv)[:]
- V = [u[:]; v[:]]
-
- # Kinetic energy and momentum of initial velocity field
- # Iteration 1 corresponds to t₀ = 0 (for unsteady simulations)
- maxdiv, umom, vmom, k = compute_conservation(V, t, setup; bc_vectors)
-
- # TODO: Maybe eps(T)^(3//4)
- if maxdiv > 1e-12
- @warn "Initial velocity field not (discretely) divergence free: $maxdiv.\n" *
- "Performing additional projection."
-
- # Make velocity field divergence free
- f = M * V + yM
- Δp = pressure_poisson(pressure_solver, f)
- V .-= 1 ./ Ω .* (G * Δp)
- end
-
- # Initial pressure: should in principle NOT be prescribed (will be calculated if p_initial)
- if isnothing(initial_pressure)
- p = pressure_additional_solve(pressure_solver, V, p, t, setup)
- else
- p .= initial_pressure.(xpp, ypp)[:]
- end
-
- V, p
-end
-
-# 3D version
-function create_initial_conditions(
- setup,
- initial_velocity_u,
- initial_velocity_v,
- initial_velocity_w,
- t;
- initial_pressure = nothing,
- pressure_solver = DirectPressureSolver(setup),
-)
- (; grid, operators) = setup
- (; xu, yu, zu, xv, yv, zv, xw, yw, zw, xpp, ypp, zpp, Ω) = grid
- (; G, M) = operators
-
- # Boundary conditions
- bc_vectors = get_bc_vectors(setup, t)
- (; yM) = bc_vectors
-
- # Allocate velocity and pressure
- u = zero(xu)[:]
- v = zero(xv)[:]
- w = zero(xw)[:]
- p = zero(xpp)[:]
-
- # Initial velocities
- u .= initial_velocity_u.(xu, yu, zu)[:]
- v .= initial_velocity_v.(xv, yv, zv)[:]
- w .= initial_velocity_w.(xw, yw, zw)[:]
- V = [u; v; w]
-
- # Kinetic energy and momentum of initial velocity field
- # Iteration 1 corresponds to t₀ = 0 (for unsteady simulations)
- maxdiv, umom, vmom, wmom, k = compute_conservation(V, t, setup; bc_vectors)
-
- # TODO: Maybe eps(T)^(3//4)
- if maxdiv > 1e-12
- @warn "Initial velocity field not (discretely) divergence free: $maxdiv.\n" *
- "Performing additional projection."
-
- # Make velocity field divergence free
- f = M * V + yM
- Δp = pressure_poisson(pressure_solver, f)
- V .-= 1 ./ Ω .* (G * Δp)
- end
-
- # Initial pressure: should in principle NOT be prescribed (will be calculated if p_initial)
- if isnothing(initial_pressure)
- p = pressure_additional_solve(pressure_solver, V, p, t, setup)
- else
- p .= initial_pressure.(xpp, ypp, zpp)[:]
- end
-
- V, p
-end
-
-function create_spectrum_2(x, A, σ, s)
- T = typeof(A)
- AT = typeof(x)
- nx, ny = size(x)
- Kx = nx ÷ 2
- Ky = ny ÷ 2
- i = AT(reshape(1:Kx, :, 1))
- j = AT(reshape(1:Ky, 1, :))
- p = T(π)
- a = @. A * (
- 1 / sqrt((2p)^2 * 2σ^2) *
- exp(-((i - s)^2 + (j - s)^2) / 2σ^2) *
- exp(-2p * im * rand(T))
- )
- [
- a reverse(a; dims = 2)
- reverse(a; dims = 1) reverse(a)
- ]
-end
-
-function create_spectrum_3(x, A, σ, s)
- T = typeof(A)
- AT = typeof(x)
- nx, ny, nz = size(x)
- Kx = nx ÷ 2
- Ky = ny ÷ 2
- Kz = nz ÷ 2
- i = AT(reshape(1:Kx, :, 1, 1))
- j = AT(reshape(1:Ky, 1, :, 1))
- k = AT(reshape(1:Kz, 1, 1, :))
- p = T(π)
- a = @. A * (
- 1 / sqrt((2p)^3 * 3σ^2) *
- exp(-((i - s)^2 + (j - s)^2 + (k - s)^2) / 2σ^2) *
- exp(-2p * im * rand(T))
- )
- [
- a reverse(a; dims = 2); reverse(a; dims = 1) reverse(a; dims = (1, 2));;;
- reverse(a; dims = 3) reverse(a; dims = (2, 3)); reverse(a; dims = (1, 3)) reverse(a)
- ]
-end
-
-"""
- random_field(
- setup, K;
- A = 1_000_000,
- σ = 30,
- s = 5,
- pressure_solver = DirectPressureSolver(setup),
- )
-
-Create random field.
-
-- `K`: Maximum wavenumber
-- `A`: Eddy amplitude
-- `σ`: Variance
-- `s` Wavenumber offset before energy starts decaying
-"""
-function random_field end
-
-random_field(setup; kwargs...) = random_field(setup.grid.dimension, setup; kwargs...)
-
-# 2D version
-function random_field(
- ::Dimension{2},
- setup;
- A = convert(eltype(setup.grid.x), 1_000_000),
- σ = convert(eltype(setup.grid.x), 30),
- s = 5,
- pressure_solver = DirectPressureSolver(setup),
-)
- (; xu, xv, Ω) = setup.grid
- (; G, M) = setup.operators
-
- T = eltype(Ω)
-
- u = real.(ifft(create_spectrum_2(xu, A, σ, s)))
- v = real.(ifft(create_spectrum_2(xv, A, σ, s)))
- V = [reshape(u, :); reshape(v, :)]
- f = M * V
- p = zero(f)
-
- # Boundary conditions
- bc_vectors = get_bc_vectors(setup, T(0))
- (; yM) = bc_vectors
-
- # Make velocity field divergence free
- f = M * V + yM
- Δp = pressure_poisson(pressure_solver, f)
- V .-= 1 ./ Ω .* (G * Δp)
- p = pressure_additional_solve(pressure_solver, V, p, T(0), setup; bc_vectors)
-
- V, p
-end
-
-# 3D version
-function random_field(
- ::Dimension{3},
- setup;
- A = convert(eltype(setup.grid.x), 1_000_000),
- σ = convert(eltype(setup.grid.x), 30),
- s = 5,
- pressure_solver = DirectPressureSolver(setup),
-)
- (; xu, xv, xw, Ω) = setup.grid
- (; G, M) = setup.operators
-
- T = eltype(Ω)
-
- u = real.(ifft(create_spectrum_3(xu, A, σ, s)))
- v = real.(ifft(create_spectrum_3(xv, A, σ, s)))
- w = real.(ifft(create_spectrum_3(xw, A, σ, s)))
- V = [reshape(u, :); reshape(v, :); reshape(w, :)]
- f = M * V
- p = zero(f)
-
- # Boundary conditions
- bc_vectors = get_bc_vectors(setup, T(0))
- (; yM) = bc_vectors
-
- # Make velocity field divergence free
- f = M * V + yM
- Δp = pressure_poisson(pressure_solver, f)
- V .-= 1 ./ Ω .* (G * Δp)
- p = pressure_additional_solve(pressure_solver, V, p, T(0), setup; bc_vectors)
-
- V, p
-end
diff --git a/src/pressure.jl b/src/pressure.jl
new file mode 100644
index 000000000..50fc406bb
--- /dev/null
+++ b/src/pressure.jl
@@ -0,0 +1,474 @@
+# Note: poisson!(...) is a wrapper around psolver!(...) since we need to define
+# a shared rrule. This could also be achieved by something like:
+# abstract type AbstractPSolver end
+# rrule(s::AbstractPSolver)(f) = (s(f), φ -> (NoTangent(), s(φ)))
+# (s::AbstractPSolver)(f) = s(zero(...), f)
+# struct PSolver < AbstractPSolver ... end
+# (::PSolver)(p, f) = # solve Poisson
+
+"""
+ poisson(psolver, f)
+
+Solve the Poisson equation for the pressure with right hand side `f` at time `t`.
+For periodic and no-slip BC, the sum of `f` should be zero.
+
+Non-mutating/allocating/out-of-place version.
+
+See also [`poisson!`](@ref).
+"""
+poisson(psolver, f) = poisson!(psolver, zero(f), f)
+
+# Laplacian is auto-adjoint
+ChainRulesCore.rrule(::typeof(poisson), psolver, f) =
+ (poisson(psolver, f), φ -> (NoTangent(), NoTangent(), poisson(psolver, unthunk(φ))))
+
+"""
+ poisson!(solver, p, f)
+
+Solve the Poisson equation for the pressure with right hand side `f` at time `t`.
+For periodic and no-slip BC, the sum of `f` should be zero.
+
+Mutating/non-allocating/in-place version.
+
+See also [`poisson`](@ref).
+"""
+poisson!(psolver, p, f) = psolver(p, f)
+
+"""
+ pressure!(p, u, temp, t, setup; psolver, F, div)
+
+Compute pressure from velocity field. This makes the pressure compatible with the velocity
+field, resulting in same order pressure as velocity.
+"""
+function pressure!(p, u, temp, t, setup; psolver, F, div)
+ (; grid) = setup
+ (; dimension, Iu, Ip, Ω) = grid
+ D = dimension()
+ momentum!(F, u, temp, t, setup)
+ apply_bc_u!(F, t, setup; dudt = true)
+ divergence!(div, F, setup)
+ @. div *= Ω
+ poisson!(psolver, p, div)
+ apply_bc_p!(p, t, setup)
+ p
+end
+
+"""
+ pressure(u, temp, t, setup; psolver)
+
+Compute pressure from velocity field. This makes the pressure compatible with the velocity
+field, resulting in same order pressure as velocity.
+"""
+function pressure(u, temp, t, setup; psolver)
+ (; grid) = setup
+ (; dimension, Iu, Ip, Ω) = grid
+ D = dimension()
+ F = momentum(u, temp, t, setup)
+ F = apply_bc_u(F, t, setup; dudt = true)
+ div = divergence(F, setup)
+ div = @. div * Ω
+ p = poisson(psolver, div)
+ p = apply_bc_p(p, t, setup)
+ p
+end
+
+"""
+ project(u, setup; psolver)
+
+Project velocity field onto divergence-free space.
+"""
+function project(u, setup; psolver)
+ (; Ω) = setup.grid
+ T = eltype(u[1])
+
+ # Divergence of tentative velocity field
+ div = divergence(u, setup)
+ div = @. div * Ω
+
+ # Solve the Poisson equation
+ p = poisson(psolver, div)
+
+ # Apply pressure correction term
+ p = apply_bc_p(p, T(0), setup)
+ G = pressuregradient(p, setup)
+ u .- G
+end
+
+"""
+ project!(u, setup; psolver, div, p)
+
+Project velocity field onto divergence-free space.
+"""
+function project!(u, setup; psolver, div, p)
+ (; Ω) = setup.grid
+ T = eltype(u[1])
+
+ # Divergence of tentative velocity field
+ divergence!(div, u, setup)
+ @. div *= Ω
+
+ # Solve the Poisson equation
+ poisson!(psolver, p, div)
+ apply_bc_p!(p, T(0), setup)
+
+ # Apply pressure correction term
+ applypressure!(u, p, setup)
+end
+
+"""
+ default_psolver(setup)
+
+Get default Poisson solver from setup.
+"""
+function default_psolver(setup)
+ (; grid, boundary_conditions) = setup
+ (; dimension, Δ) = grid
+ D = dimension()
+ Δx = first.(Array.(Δ))
+ isperiodic =
+ all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions)
+ isuniform = all(α -> all(≈(Δx[α]), Δ[α]), 1:D)
+ if isperiodic && isuniform
+ psolver_spectral(setup)
+ else
+ psolver_direct(setup)
+ end
+end
+
+"""
+ poisson_direct(setup)
+
+Create direct Poisson solver using an appropriate matrix decomposition.
+"""
+psolver_direct(setup) = psolver_direct(setup.grid.x[1], setup) # Dispatch on array type
+
+# CPU version
+function psolver_direct(::Array, setup)
+ (; grid, boundary_conditions) = setup
+ (; x, Np, Ip) = grid
+ T = Float64 # This is currently required for SuiteSparse
+ L = laplacian_mat(setup)
+ isdefinite =
+ any(bc -> bc[1] isa PressureBC || bc[2] isa PressureBC, boundary_conditions)
+ if isdefinite
+ # No extra DOF
+ ftemp = zeros(T, prod(Np))
+ ptemp = zeros(T, prod(Np))
+ viewrange = (:)
+ else
+ # With extra DOF
+ ftemp = zeros(T, prod(Np) + 1)
+ ptemp = zeros(T, prod(Np) + 1)
+ e = ones(T, size(L, 2))
+ L = [L e; e' 0]
+ viewrange = 1:prod(Np)
+ end
+ fact = factorize(L)
+ function psolve!(p, f)
+ copyto!(view(ftemp, viewrange), view(view(f, Ip), :))
+ ptemp .= fact \ ftemp
+ copyto!(view(view(p, Ip), :), eltype(p).(view(ptemp, viewrange)))
+ p
+ end
+end
+
+"""
+ psolver_cg_matrix(setup; kwargs...)
+
+Conjugate gradients iterative Poisson solver.
+The `kwargs` are passed to the `cg!` function
+from IterativeSolvers.jl.
+"""
+function psolver_cg_matrix(setup; kwargs...)
+ (; x, Np, Ip) = grid
+ L = laplacian_mat(setup)
+ isdefinite =
+ any(bc -> bc[1] isa PressureBC || bc[2] isa PressureBC, boundary_conditions)
+ if isdefinite
+ # No extra DOF
+ ftemp = fill!(similar(x[1], prod(Np)), 0)
+ ptemp = fill!(similar(x[1], prod(Np)), 0)
+ viewrange = (:)
+ else
+ # With extra DOF
+ ftemp = fill!(similar(x[1], prod(Np) + 1), 0)
+ ptemp = fill!(similar(x[1], prod(Np) + 1), 0)
+ e = fill!(similar(x[1], prod(Np)), 1)
+ L = [L e; e' 0]
+ viewrange = 1:prod(Np)
+ end
+ function psolve!(p, f)
+ copyto!(view(ftemp, viewrange), view(view(f, Ip), :))
+ cg!(ptemp, L, ftemp; kwargs...)
+ copyto!(view(view(p, Ip), :), view(ptemp, viewrange))
+ p
+ end
+end
+
+# Preconditioner
+function create_laplace_diag(setup)
+ (; grid, workgroupsize) = setup
+ (; dimension, Δ, Δu, N, Np, Ip, Ω) = grid
+ D = dimension()
+ δ = Offset{D}()
+ @kernel function _laplace_diag!(z, p, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ d = zero(eltype(z))
+ for α = 1:length(I)
+ d -= Ω[I] / Δ[α][I[α]] * (1 / Δu[α][I[α]] + 1 / Δu[α][I[α]-1])
+ end
+ z[I] = -p[I] / d
+ end
+ ndrange = Np
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ laplace_diag(z, p) = _laplace_diag!(get_backend(z), workgroupsize)(z, p, I0; ndrange)
+end
+
+"""
+ psolver_cg(
+ setup;
+ abstol = zero(eltype(setup.grid.x[1])),
+ reltol = sqrt(eps(eltype(setup.grid.x[1]))),
+ maxiter = prod(setup.grid.Np),
+ preconditioner = create_laplace_diag(setup),
+ )
+
+Conjugate gradients iterative Poisson solver.
+"""
+function psolver_cg(
+ setup;
+ abstol = zero(eltype(setup.grid.x[1])),
+ reltol = sqrt(eps(eltype(setup.grid.x[1]))),
+ maxiter = prod(setup.grid.Np),
+ preconditioner = create_laplace_diag(setup),
+)
+ (; grid, workgroupsize) = setup
+ (; Np, Ip, Ω) = grid
+ T = eltype(setup.grid.x[1])
+ r = similar(setup.grid.x[1], setup.grid.N)
+ L = similar(setup.grid.x[1], setup.grid.N)
+ q = similar(setup.grid.x[1], setup.grid.N)
+ function psolve!(p, f)
+ function innerdot(a, b)
+ @kernel function innerdot!(d, a, b, I0)
+ I = @index(Global, Cartesian)
+ I = I + I0
+ d[I-I+I0] += a[I] * b[I]
+ # a[I] = b[I]
+ end
+ # d = zero(eltype(a))
+ I0 = first(Ip)
+ I0 -= oneunit(I0)
+ d = fill!(similar(a, ntuple(Returns(1), length(I0))), 0),
+ innerdot!(get_backend(a), workgroupsize)(d, a, b, I0; ndrange = Np)
+ d[]
+ end
+
+ p .= 0
+
+ # Initial residual
+ laplacian!(L, p, setup)
+
+ # Initialize
+ q .= 0
+ r .= f .- L
+ ρ_prev = one(T)
+ # residual = norm(r[Ip])
+ residual = sqrt(sum(abs2, view(r, Ip)))
+ # residual = norm(r)
+ tolerance = max(reltol * residual, abstol)
+ iteration = 0
+
+ while iteration < maxiter && residual > tolerance
+ preconditioner(L, r)
+
+ # ρ = sum(L[Ip] .* r[Ip])
+ ρ = dot(view(L, Ip), view(r, Ip))
+ # ρ = innerdot(L, r)
+ # ρ = dot(L, r)
+
+ β = ρ / ρ_prev
+ q .= L .+ β .* q
+
+ # Periodic/symmetric padding (maybe)
+ apply_bc_p!(q, T(0), setup)
+ laplacian!(L, q, setup)
+ # α = ρ / sum(q[Ip] .* L[Ip])
+ α = ρ / dot(view(q, Ip), view(L, Ip))
+ # α = ρ / innerdot(q, L)
+ # α = ρ / dot(q, L)
+
+ p .+= α .* q
+ r .-= α .* L
+
+ ρ_prev = ρ
+ # residual = norm(r[Ip])
+ residual = sqrt(sum(abs2, view(r, Ip)))
+ # residual = sqrt(sum(abs2, r))
+ # residual = sqrt(innerdot(r, r))
+
+ iteration += 1
+ end
+
+ # @show iteration residual tolerance
+
+ p
+ end
+end
+
+"""
+ psolver_spectral(setup)
+
+Create spectral Poisson solver from setup.
+"""
+function psolver_spectral(setup)
+ (; grid, boundary_conditions) = setup
+ (; dimension, Δ, Np, Ip, x) = grid
+
+ D = dimension()
+ T = eltype(Δ[1])
+
+ Δx = first.(Array.(Δ))
+
+ @assert(
+ all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions),
+ "Spectral psolver only implemented for periodic boundary conditions",
+ )
+
+ @assert(
+ all(α -> all(≈(Δx[α]), Δ[α]), 1:D),
+ "Spectral psolver requires uniform grid along each dimension",
+ )
+
+ # Fourier transform of the discretization
+ # Assuming uniform grid, although Δx[1] and Δx[2] do not need to be the same
+
+ k = ntuple(
+ d -> reshape(
+ 0:Np[d]-1,
+ ntuple(Returns(1), d - 1)...,
+ :,
+ ntuple(Returns(1), D - d)...,
+ ),
+ D,
+ )
+
+ Ahat = fill!(similar(x[1], Complex{T}, Np), 0)
+ Tπ = T(π) # CUDA doesn't like pi
+ for d = 1:D
+ @. Ahat += sin(k[d] * Tπ / Np[d])^2 / Δx[d]^2
+ end
+
+ # Scale with Δx*Δy*Δz, since we solve the PDE in integrated form
+ Ahat .*= 4 * prod(Δx)
+
+ # Pressure is determined up to constant. By setting the constant
+ # scaling factor to 1, we preserve the average.
+ # Note use of singleton range 1:1 instead of scalar index 1
+ # (otherwise CUDA gets annoyed)
+ Ahat[1:1] .= 1
+
+ # Placeholders for intermediate results
+ phat = zero(Ahat)
+ fhat = zero(Ahat)
+ plan = plan_fft(fhat)
+
+ function psolver(p, f)
+ f = view(f, Ip)
+
+ fhat .= complex.(f)
+
+ # Fourier transform of right hand side
+ mul!(phat, plan, fhat)
+
+ # Solve for coefficients in Fourier space
+ @. fhat = -phat / Ahat
+
+ # Pressure is determined up to constant. We set this to 0 (instead of
+ # phat[1] / 0 = Inf)
+ # Note use of singleton range 1:1 instead of scalar index 1
+ # (otherwise CUDA gets annoyed)
+ fhat[1:1] .= 0
+
+ # Transform back
+ ldiv!(phat, plan, fhat)
+ @. p[Ip] = real(phat)
+
+ p
+ end
+end
+
+"""
+ psolver_spectral_lowmemory(setup)
+
+Create spectral Poisson solver from setup.
+This one is slower than `psolver_spectral` but occupies less memory.
+"""
+function psolver_spectral_lowmemory(setup)
+ (; grid, boundary_conditions) = setup
+ (; dimension, Δ, Np, Ip, x) = grid
+
+ D = dimension()
+ T = eltype(Δ[1])
+
+ Δx = Array(Δ[1])[1]
+
+ @assert(
+ all(bc -> bc[1] isa PeriodicBC && bc[2] isa PeriodicBC, boundary_conditions),
+ "Spectral psolver only implemented for periodic boundary conditions",
+ )
+
+ @assert(
+ all(α -> all(≈(Δx), Δ[α]), 1:D),
+ "Spectral psolver requires uniform grid along each dimension",
+ )
+
+ @assert all(n -> n == Np[1], Np)
+
+ # Fourier transform of the discretization
+ # Assuming uniform grid, although Δx[1] and Δx[2] do not need to be the same
+
+ k = 0:Np[1]-1
+
+ Tπ = T(π) # CUDA doesn't like pi
+ ahat = fill!(similar(x[1], Np[1]), 0)
+ @. ahat = 4 * Δx^D * sin(k * Tπ / Np[1])^2 / Δx^2
+
+ # Placeholders for intermediate results
+ phat = fill!(similar(x[1], Complex{T}, Np), 0)
+
+ function psolve!(p, f)
+ f = view(f, Ip)
+
+ phat .= complex.(f)
+
+ # Fourier transform of right hand side
+ fft!(phat)
+
+ # Solve for coefficients in Fourier space
+ if D == 2
+ ax = ahat
+ ay = reshape(ahat, 1, :)
+ @. phat = -phat / (ax + ay)
+ else
+ ax = ahat
+ ay = reshape(ahat, 1, :)
+ az = reshape(ahat, 1, 1, :)
+ @. phat = -phat / (ax + ay + az)
+ end
+
+ # Pressure is determined up to constant. We set this to 0 (instead of
+ # phat[1] / 0 = Inf)
+ # Note use of singleton range 1:1 instead of scalar index 1
+ # (otherwise CUDA gets annoyed)
+ phat[1:1] .= 0
+
+ # Transform back
+ ifft!(phat)
+ @. p[Ip] = real(phat)
+
+ p
+ end
+end
diff --git a/src/processors.jl b/src/processors.jl
new file mode 100644
index 000000000..124913d76
--- /dev/null
+++ b/src/processors.jl
@@ -0,0 +1,651 @@
+"""
+ processor(initialize, finalize = (initialized, state) -> initialized)
+
+Process results from time stepping. Before time stepping, the `initialize`
+function is called on an observable of the time stepper `state`, returning
+`initialized`. The observable is updated every time step.
+
+After timestepping, the `finalize` function is called on `initialized` and the
+final `state`.
+
+See the following example:
+
+```example
+function initialize(state)
+ s = 0
+ println("Let's sum up the time steps")
+ on(state) do (; n, t)
+ println("The summand is \$n, the time is \$t")
+ s = s + n
+ end
+ s
+end
+
+finalize(i, state) = println("The final sum (at time t=\$(state.t)) is \$s")
+p = processor(initialize, finalize)
+```
+
+When solved for 6 time steps from t=0 to t=2 the displayed output is
+
+```
+Let's sum up the time steps
+The summand is 0, the time is 0.0
+The summand is 1, the time is 0.4
+The summand is 2, the time is 0.8
+The summand is 3, the time is 1.2
+The summand is 4, the time is 1.6
+The summand is 5, the time is 2.0
+The final sum (at time t=2.0) is 15
+```
+"""
+processor(initialize, finalize = (initialized, state) -> initialized) =
+ (; initialize, finalize)
+
+"""
+ timelogger(; nupdate = 1)
+
+Create processor that logs time step information.
+"""
+timelogger(; nupdate = 1) =
+ processor() do state
+ on(state) do (; t, n)
+ n % nupdate == 0 || return
+ @printf "Iteration %d\tt = %g\n" n t
+ end
+ nothing
+ end
+
+"""
+ vtk_writer(;
+ setup,
+ nupdate = 1,
+ dir = "output",
+ filename = "solution",
+ fields = (:velocity, :pressure, :vorticity),
+ )
+
+Create processor that writes the solution every `nupdate` time steps to a VTK
+file. The resulting Paraview data collection file is stored in
+`"\$dir/\$filename.pvd"`.
+"""
+vtk_writer(;
+ setup,
+ nupdate = 1,
+ dir = "output",
+ filename = "solution",
+ fields = (:velocity, :pressure, :vorticity),
+ psolver = nothing,
+) =
+ processor((pvd, state) -> vtk_save(pvd)) do state
+ (; grid) = setup
+ (; dimension, xp) = grid
+ D = dimension()
+ ispath(dir) || mkpath(dir)
+ pvd = paraview_collection(joinpath(dir, filename))
+ xparr = Array.(xp)
+ (; u) = state[]
+ if :velocity ∈ fields
+ up = interpolate_u_p(u, setup)
+ uparr = Array.(up)
+ # ParaView prefers 3D vectors. Add zero z-component.
+ D == 2 && (uparr = (uparr..., zero(up[1])))
+ end
+ if :pressure ∈ fields
+ if isnothing(psolver)
+ @info "Creating new pressure solver for vtk_writer"
+ psolver = default_psolver(setup)
+ end
+ F = zero.(u)
+ div = zero(u[1])
+ p = zero(u[1])
+ parr = adapt(Array, p)
+ end
+ if :vorticity ∈ fields
+ ω = vorticity(u, setup)
+ ωp = interpolate_ω_p(ω, setup)
+ ωparr = adapt(Array, ωp)
+ end
+ on(state) do (; u, t, n)
+ n % nupdate == 0 || return
+ tformat = replace(string(t), "." => "p")
+ vtk_grid("$(dir)/$(filename)_t=$tformat", xparr...) do vtk
+ if :velocity ∈ fields
+ interpolate_u_p!(up, u, setup)
+ copyto!.(uparr, up)
+ vtk["velocity"] = uparr
+ end
+ if :pressure ∈ fields
+ pressure!(p, u, temp, t, setup; psolver, F, div)
+ vtk["pressure"] = copyto!(parr, p)
+ end
+ if :vorticity ∈ fields
+ vorticity!(ω, u, setup)
+ interpolate_ω_p!(ωp, ω, setup)
+ D == 2 ? copyto!(ωparr, ωp) : copyto!.(ωparr, ωp)
+ vtk["vorticity"] = ωparr
+ end
+ pvd[t] = vtk
+ end
+ end
+ pvd
+ end
+
+"""
+ fieldsaver(; setup, nupdate = 1)
+
+Create processor that stores the solution and time every `nupdate` time step.
+"""
+fieldsaver(; setup, nupdate = 1) =
+ processor() do state
+ T = eltype(setup.grid.x[1])
+ (; u) = state[]
+ fields = (; u = fill(Array.(u), 0), t = zeros(T, 0))
+ on(state) do (; u, p, t, n)
+ n % nupdate == 0 || return
+ push!(fields.u, Array.(u))
+ push!(fields.t, t)
+ end
+ fields
+ end
+
+"""
+ animator(; setup, path, plot = fieldplot, nupdate = 1, kwargs...)
+
+Animate a plot of the solution every `update` iteration.
+The animation is saved to `path`, which should have one
+of the following extensions:
+
+- ".mkv"
+- ".mp4"
+- ".webm"
+- ".gif"
+
+The plot is determined by a `plotter` processor.
+Additional `kwargs` are passed to `plot`.
+"""
+animator(;
+ setup,
+ path,
+ plot = fieldplot,
+ nupdate = 1,
+ framerate = 24,
+ visible = true,
+ kwargs...,
+) =
+ processor((stream, state) -> save(path, stream)) do outerstate
+ ispath(dirname(path)) || mkpath(dirname(path))
+ state = Observable(outerstate[])
+ fig = plot(state; setup, kwargs...)
+ visible && display(fig)
+ stream = VideoStream(fig; framerate, visible)
+ on(outerstate) do outerstate
+ outerstate.n % nupdate == 0 || return
+ state[] = outerstate
+ recordframe!(stream)
+ end
+ stream
+ end
+
+"""
+ realtimeplotter(;
+ setup,
+ plot = fieldplot,
+ nupdate = 1,
+ displayfig = true,
+ displayupdates = false,
+ sleeptime = nothing,
+ kwargs...,
+ )
+
+Processor for plotting the solution every `nupdate` time step.
+
+The `sleeptime` is slept at every update, to give Makie time to update the
+plot. Set this to `nothing` to skip sleeping.
+
+Additional `kwargs` are passed to the `plot` function.
+"""
+realtimeplotter(;
+ setup,
+ plot = fieldplot,
+ nupdate = 1,
+ displayfig = true,
+ displayupdates = false,
+ sleeptime = nothing,
+ kwargs...,
+) =
+ processor() do outerstate
+ state = Observable(outerstate[])
+ fig = plot(state; setup, kwargs...)
+ displayfig && display(fig)
+ on(outerstate) do outerstate
+ outerstate.n % nupdate == 0 || return
+ state[] = outerstate
+ displayupdates && display(fig)
+ isnothing(sleeptime) || sleep(sleeptime)
+ end
+ fig
+ end
+
+"""
+ fieldplot(
+ state;
+ setup,
+ fieldname = :vorticity,
+ type = nothing,
+ kwargs...,
+ )
+
+Plot `state` field in pressure points.
+If `state` is `Observable`, then the plot is interactive.
+
+Available fieldnames are:
+
+- `:velocity`,
+- `:vorticity`,
+- `:streamfunction`,
+- `:pressure`.
+
+Available plot `type`s for 2D are:
+
+- `heatmap` (default),
+- `image`,
+- `contour`,
+- `contourf`.
+
+Available plot `type`s for 3D are:
+
+- `contour` (default).
+
+The `alpha` value gets passed to `contour` in 3D.
+"""
+fieldplot(state; setup, kwargs...) = fieldplot(
+ setup.grid.dimension,
+ state isa Observable ? state : Observable(state);
+ setup,
+ kwargs...,
+)
+
+function fieldplot(
+ ::Dimension{2},
+ state;
+ setup,
+ fieldname = :vorticity,
+ psolver = nothing,
+ type = heatmap,
+ equal_axis = true,
+ docolorbar = true,
+ size = nothing,
+ title = nothing,
+ kwargs...,
+)
+ (; boundary_conditions, grid) = setup
+ (; dimension, xlims, x, xp, Ip, Δ) = grid
+ D = dimension()
+
+ xf = Array.(getindex.(setup.grid.xp, Ip.indices))
+
+ (; u, temp, t) = state[]
+ _f = if fieldname in (1, 2)
+ up = interpolate_u_p(u, setup)
+ up[fieldname]
+ elseif fieldname == :velocity
+ up = interpolate_u_p(u, setup)
+ upnorm = zero(up[1])
+ elseif fieldname == :vorticity
+ ω = vorticity(u, setup)
+ ωp = interpolate_ω_p(ω, setup)
+ elseif fieldname == :streamfunction
+ ψ = get_streamfunction(setup, u, t)
+ elseif fieldname == :pressure
+ if isnothing(psolver)
+ @info "Creating new pressure solver for fieldplot"
+ psolver = default_psolver(setup)
+ end
+ F = zero.(u)
+ div = zero(u[1])
+ p = zero(u[1])
+ elseif fieldname == :V1
+ B, V = tensorbasis(u, setup)
+ V[1]
+ elseif fieldname == :V2
+ B, V = tensorbasis(u, setup)
+ V[2]
+ elseif fieldname == :temperature
+ temp
+ end
+ _f = Array(_f)[Ip]
+ field = lift(state) do (; u, temp, t)
+ f = if fieldname in (1, 2)
+ interpolate_u_p!(up, u, setup)
+ up[fieldname]
+ elseif fieldname == :velocity
+ interpolate_u_p!(up, u, setup)
+ map((u, v) -> √sum(u^2 + v^2), up...)
+ @. upnorm = sqrt(up[1]^2 + up[2]^2)
+ elseif fieldname == :vorticity
+ apply_bc_u!(u, t, setup)
+ vorticity!(ω, u, setup)
+ interpolate_ω_p!(ωp, ω, setup)
+ elseif fieldname == :streamfunction
+ get_streamfunction!(setup, ψ, u, t)
+ elseif fieldname == :pressure
+ pressure!(p, u, temp, t, setup; psolver, F, div)
+ elseif fieldname == :V1
+ tensorbasis!(B, V, u, setup)
+ V[1]
+ elseif fieldname == :V2
+ tensorbasis!(B, V, u, setup)
+ -V[2]
+ elseif fieldname == :temperature
+ temp
+ end
+ # Array(f)[Ip]
+ copyto!(_f, view(f, Ip))
+ end
+
+ lims = lift(field) do f
+ if type ∈ (heatmap, image)
+ lims = get_lims(f)
+ elseif type ∈ (contour, contourf)
+ if ≈(extrema(f)...; rtol = 1e-10)
+ μ = mean(f)
+ a = μ - 1
+ b = μ + 1
+ f[1] += 1
+ f[end] -= 1
+ else
+ a, b = get_lims(f)
+ end
+ lims = (a, b)
+ end
+ lims
+ end
+
+ if type ∈ (heatmap, image)
+ kwargs = (; colorrange = lims, kwargs...)
+ elseif type ∈ (contour, contourf)
+ kwargs = (;
+ extendlow = :auto,
+ extendhigh = :auto,
+ levels = @lift(LinRange($(lims)..., 10)),
+ colorrange = lims,
+ kwargs...,
+ )
+ end
+
+ axis = (;
+ xlabel = "x",
+ ylabel = "y",
+ title = isnothing(title) ? titlecase(string(fieldname)) : title,
+ limits = (xlims[1]..., xlims[2]...),
+ )
+ equal_axis && (axis = (axis..., aspect = DataAspect()))
+
+ # Image requires boundary coordinates only
+ if type == image
+ Δx = first.(Array.(Δ))
+ @assert all(≈(Δx[1]), Δx) "Image requires rectangular pixels"
+ @assert(all(α -> all(≈(Δx[α]), Δ[α]), 1:D), "Image requires uniform grid",)
+ xf = map(extrema, xf)
+ end
+
+ size = isnothing(size) ? (;) : (; size)
+ fig = Figure(; size...)
+ ax, hm = type(fig[1, 1], xf..., field; axis, kwargs...)
+ docolorbar && Colorbar(fig[1, 2], hm)
+
+ fig
+end
+
+function fieldplot(
+ ::Dimension{3},
+ state;
+ setup,
+ psolver = nothing,
+ fieldname = :eig2field,
+ alpha = convert(eltype(setup.grid.x[1]), 0.1),
+ isorange = convert(eltype(setup.grid.x[1]), 0.5),
+ equal_axis = true,
+ levels = LinRange{eltype(setup.grid.x[1])}(-10, 5, 10),
+ docolorbar = false,
+ size = nothing,
+ logtol = eps(setup.T),
+ kwargs...,
+)
+ (; boundary_conditions, grid) = setup
+ (; xlims, x, xp, Ip) = grid
+
+ xf = Array.(getindex.(setup.grid.xp, Ip.indices))
+ dxf = diff.(xf)
+ if all(α -> all(≈(dxf[α][1]), dxf[α]), 1:3)
+ xf = ntuple(α -> LinRange(xf[α][1], xf[α][end], length(xf[α])), 3)
+ end
+
+ (; u) = state[]
+ T = eltype(xf[1])
+ if fieldname == :velocity
+ elseif fieldname == :vorticity
+ elseif fieldname == :streamfunction
+ elseif fieldname == :pressure
+ if isnothing(psolver)
+ @info "Creating new pressure solver for fieldplot"
+ psolver = default_psolver(setup)
+ end
+ F = zero.(u)
+ div = zero(u[1])
+ p = zero(u[1])
+ elseif fieldname == :Dfield
+ if isnothing(psolver)
+ @info "Creating new pressure solver for fieldplot"
+ psolver = default_psolver(setup)
+ end
+ F = zero.(u)
+ div = zero(u[1])
+ p = zero(u[1])
+ G = similar.(u)
+ d = similar(u[1])
+ elseif fieldname == :Qfield
+ Q = similar(u[1])
+ elseif fieldname == :eig2field
+ λ = similar(u[1])
+ elseif fieldname in union(Symbol.(["B$i" for i = 1:11]), Symbol.(["V$i" for i = 1:5]))
+ sym = string(fieldname)[1]
+ sym = sym == 'B' ? 1 : 2
+ idx = parse(Int, string(fieldname)[2:end])
+ tb = tensorbasis(u, setup)
+ tb[sym][idx]
+ elseif fieldname == :temperature
+ else
+ error("Unknown fieldname")
+ end
+
+ field = lift(state) do (; u, temp, t)
+ f = if fieldname == :velocity
+ up = interpolate_u_p(u, setup)
+ map((u, v, w) -> √sum(u^2 + v^2 + w^2), up...)
+ elseif fieldname == :vorticity
+ ωp = interpolate_ω_p(vorticity(u, setup), setup)
+ map((u, v, w) -> √sum(u^2 + v^2 + w^2), ωp...)
+ elseif fieldname == :streamfunction
+ get_streamfunction(setup, u, t)
+ elseif fieldname == :pressure
+ pressure!(p, u, temp, t, setup; psolver, F, div)
+ elseif fieldname == :Dfield
+ pressure!(p, u, temp, t, setup; psolver, F, div)
+ Dfield!(d, G, p, setup)
+ din = view(d, Ip)
+ @. din = log(max(logtol, din))
+ d
+ elseif fieldname == :Qfield
+ Qfield!(Q, u, setup)
+ Qin = view(Q, Ip)
+ @. Qin = log(max(logtol, Qin))
+ Q
+ elseif fieldname == :eig2field
+ eig2field!(λ, u, setup)
+ λin = view(λ, Ip)
+ @. λin .= log(max(logtol, -λin))
+ λ
+ elseif fieldname in
+ union(Symbol.(["B$i" for i = 1:11]), Symbol.(["V$i" for i = 1:5]))
+ tensorbasis!(tb..., u, setup)
+ tb[sym][idx]
+ elseif fieldname == :temperature
+ temp
+ end
+ Array(f)[Ip]
+ end
+
+ # color = lift(state) do (; temp)
+ # Array(view(temp, Ip))
+ # end
+ # colorrange = lift(state) do (; temp)
+ # extrema(view(temp, Ip))
+ # end
+
+ # lims = @lift get_lims($field)
+ lims = isnothing(levels) ? lift(get_lims, field) : extrema(levels)
+
+ isnothing(levels) && (levels = @lift(LinRange($(lims)..., 10)))
+
+ # aspect = equal_axis ? (; aspect = :data) : (;)
+ size = isnothing(size) ? (;) : (; size)
+ fig = Figure(; size...)
+ # ax = Axis3(fig[1, 1]; title = titlecase(string(fieldname)), aspect...)
+ hm = contour(
+ fig[1, 1],
+ # ax,
+ xf...,
+ field;
+ levels,
+ # color,
+ # colorrange,
+ colorrange = lims,
+ # colorrange = extrema(levels),
+ alpha,
+ isorange,
+ # highclip = :red,
+ # lowclip = :red,
+ kwargs...,
+ )
+ docolorbar && Colorbar(fig[1, 2], hm)
+ fig
+end
+
+"""
+ energy_history_plot(state; setup)
+
+Create energy history plot.
+"""
+function energy_history_plot(state; setup)
+ @assert state isa Observable "Energy history requires observable state."
+ (; Ω, Ip) = setup.grid
+ e = zero(state[].u[1])
+ _points = Point2f[]
+ points = lift(state) do (; u, t)
+ kinetic_energy!(e, u, setup)
+ e .*= Ω
+ E = sum(e[Ip])
+ push!(_points, Point2f(t, E))
+ end
+ fig = lines(points; axis = (; xlabel = "t", ylabel = "Kinetic energy"))
+ on(_ -> autolimits!(fig.axis), points)
+ fig
+end
+
+"""
+ energy_spectrum_plot(state; setup)
+
+Create energy spectrum plot.
+The energy at a scalar wavenumber level ``\\kappa \\in \\mathbb{N}`` is defined by
+
+```math
+\\hat{e}(\\kappa) = \\int_{\\kappa \\leq \\| k \\|_2 < \\kappa + 1} | \\hat{e}(k) | \\mathrm{d} k,
+```
+
+as in San and Staples [San2012](@cite).
+"""
+function energy_spectrum_plot(
+ state;
+ setup,
+ npoint = 100,
+ a = typeof(setup.Re)(1 + sqrt(5)) / 2,
+)
+ state isa Observable || (state = Observable(state))
+
+ (; dimension, xp, Ip) = setup.grid
+ T = eltype(xp[1])
+ D = dimension()
+
+ (; A, κ, K) = spectral_stuff(setup; npoint, a)
+ # (; masks, κ, K) = get_spectrum(setup; npoint, a) # Mask
+ kmax = maximum(κ)
+
+ # Energy
+ # up = interpolate_u_p(state[].u, setup)
+ ehat = lift(state) do (; u, t)
+ # interpolate_u_p!(up, u, setup)
+ up = u
+ e = sum(up) do u
+ u = u[Ip]
+ uhat = fft(u)[ntuple(α -> 1:K[α], D)...]
+ # uhat = fft(u)[ntuple(α -> 1:K, D)...] # Mask
+ abs2.(uhat) ./ (2 * prod(size(uhat))^2)
+ end
+ e = A * reshape(e, :)
+ # e = [sum(e[m]) for m in masks] # Mask
+ e = max.(e, eps(T)) # Avoid log(0)
+ Array(e)
+ end
+
+ # Build inertial slope above energy
+ # krange = LinRange(1, kmax, 100)
+ # krange = collect(1, kmax)
+ # krange = [cbrt(T(kmax)), T(kmax)]
+ krange = [kmax^T(0.3), kmax^(T(0.8))]
+ # krange = [T(kmax)^(T(2) / 3), T(kmax)]
+ slope, slopelabel = D == 2 ? (-T(3), L"$k^{-3}") : (-T(5 / 3), L"$k^{-5/3}")
+ inertia = lift(ehat) do ehat
+ slopeconst = maximum(ehat ./ κ .^ slope)
+ 2 .* slopeconst .* krange .^ slope
+ end
+
+ # Nice ticks
+ logmax = round(Int, log2(kmax + 1))
+ xticks = T(2) .^ (0:logmax)
+
+ fig = Figure()
+ ax = Axis(
+ fig[1, 1];
+ xticks,
+ xlabel = "k",
+ # ylabel = "E(k)",
+ xscale = log10,
+ yscale = log10,
+ limits = (1, kmax, T(1e-8), T(1)),
+ )
+ lines!(ax, κ, ehat; label = "Kinetic energy")
+ lines!(ax, krange, inertia; label = slopelabel, linestyle = :dash)
+ axislegend(ax)
+ # autolimits!(ax)
+ on(e -> autolimits!(ax), ehat)
+ autolimits!(ax)
+ fig
+end
+
+# # Make sure the figure is fully rendered before allowing code to continue
+# if displayfig
+# render = display(espec)
+# done_rendering = Ref(false)
+# on(render.render_tic) do _
+# done_rendering[] = true
+# end
+# on(state) do s
+# # State is updated, block code execution until GLMakie has rendered
+# # figure update
+# done_rendering[] = false
+# while !done_rendering[]
+# sleep(checktime)
+# end
+# end
+# end
diff --git a/src/processors/animator.jl b/src/processors/animator.jl
deleted file mode 100644
index 614690762..000000000
--- a/src/processors/animator.jl
+++ /dev/null
@@ -1,29 +0,0 @@
-"""
- animator(setup, path; nupdate = 1, plotter = field_plotter(setup); kwargs...)
-
-Animate a plot of the solution every `update` iteration.
-The animation is saved to `path`, which should have one
-of the following extensions:
-
-- ".mkv"
-- ".mp4"
-- ".webm"
-- ".gif"
-
-The plot is determined by a `plotter` processsor.
-Addtional `kwargs` are passed to Makie's `VideoStream`.
-"""
-animator(setup, path; nupdate = 1, plotter = field_plotter(setup), kwargs...) = processor(
- function (state)
- _state = Observable(state[])
- fig = plotter.initialize(_state)
- stream = VideoStream(fig, kwargs...)
- @lift begin
- _state[] = $state
- recordframe!(stream)
- end
- stream
- end;
- finalize = (stream, step_observer) -> save(path, stream),
- nupdate,
-)
diff --git a/src/processors/processors.jl b/src/processors/processors.jl
deleted file mode 100644
index a3c48f539..000000000
--- a/src/processors/processors.jl
+++ /dev/null
@@ -1,120 +0,0 @@
-raw"""
- processor(
- initialize;
- finalize = (initialized, stepper) -> initialized,
- nupdate = 1
- )
-
-Process results from time stepping. Before time stepping, the `initialize` function is called on an observable of the time `stepper`, returning `initialized`. The observable is updated after every `nupdate` time step, triggering updates where `@lift` is used inside `initialize`.
-
-After timestepping, the `finalize` function is called on `initialized` and the final stepper.
-
-See the following example:
-
-```example
-function initialize(step_observer)
- s = 0
- println("Let's sum up the time steps")
- @lift begin
- (; n) = $step_observer
- println("The summand is $n")
- s = s + n
- end
- s
-end
-
-finalize(s, stepper) = println("The final sum (at time t=$(stepper.t)) is $s")
-p = Processor(intialize; finalize, nupdate = 5)
-```
-
-When solved for 20 time steps from t=0 to t=2 the displayed output is
-
-```
-Let's sum up the time steps
-The summand is 0
-The summand is 5
-The summand is 10
-The summand is 15
-The summand is 20
-The final sum (at time t=2.0) is 50
-```
-"""
-processor(initialize; finalize = (initialized, stepper) -> initialized, nupdate = 1) =
- (; initialize, finalize, nupdate)
-
-"""
- step_logger(; nupdate = 1)
-
-Create processor that logs time step information.
-"""
-step_logger(; nupdate = 1) = processor(function (step_observer)
- @lift begin
- (; t, n) = $step_observer
- @printf "Iteration %d\tt = %g\n" n t
- end
- nothing
-end; nupdate)
-
-"""
- vtk_writer(setup; nupdate, dir = "output", filename = "solution")
-
-Create processor that writes the solution every `nupdate` time steps to a VTK file. The resulting Paraview data
-collection file is stored in `"\$dir/\$filename.pvd"`.
-"""
-vtk_writer(setup; nupdate = 1, dir = "output", filename = "solution") = processor(
- function (step_observer)
- ispath(dir) || mkpath(dir)
- pvd = paraview_collection(joinpath(dir, filename))
- @lift begin
- (; grid) = setup
- (; dimension, xp, yp) = grid
- (; V, p, t) = $step_observer
-
- V = Array(V)
- p = Array(p)
-
- N = dimension()
- if N == 2
- coords = (xp, yp)
- elseif N == 3
- (; zp) = grid
- coords = (xp, yp, zp)
- end
-
- tformat = replace(string(t), "." => "p")
- vtk_grid("$(dir)/$(filename)_t=$tformat", coords...) do vtk
- vels = get_velocity(setup, V, t)
- if N == 2
- # ParaView prefers 3D vectors. Add zero z-component.
- wp = zeros(size(vels[1]))
- vels = (vels..., wp)
- end
- vtk["velocity"] = vels
- vtk["pressure"] = p
- pvd[t] = vtk
- end
- end
- pvd
- end;
- finalize = (pvd, step_observer) -> vtk_save(pvd),
- nupdate,
-)
-
-"""
- field_saver(setup; nupdate = 1)
-
-Create processor that stores the solution every `nupdate` time step to the vector of vectors `V` and `p`. The solution times are stored in the vector `t`.
-"""
-field_saver(setup; nupdate = 1) = processor(function (step_observer)
- T = eltype(setup.grid.x)
- _V = fill(zeros(T, 0), 0)
- _p = fill(zeros(T, 0), 0)
- _t = fill(zero(T), 0)
- @lift begin
- (; V, p, t) = $step_observer
- push!(_V, Array(V))
- push!(_p, Array(p))
- push!(_t, t)
- end
- (; V = _V, p = _p, t = _t)
-end; nupdate)
diff --git a/src/processors/real_time_plot.jl b/src/processors/real_time_plot.jl
deleted file mode 100644
index 563d0566d..000000000
--- a/src/processors/real_time_plot.jl
+++ /dev/null
@@ -1,311 +0,0 @@
-"""
- field_plotter(
- setup;
- fieldname = :vorticity,
- type = nothing,
- sleeptime = 0.001,
- alpha = 0.05,
- )
-
-Plot the solution every time the state `o` is updated.
-
-The `sleeptime` is slept at every update, to give Makie time to update the
-plot. Set this to `nothing` to skip sleeping.
-
-Available fieldnames are:
-
-- `:velocity`,
-- `:vorticity`,
-- `:streamfunction`,
-- `:pressure`.
-
-Available plot `type`s for 2D are:
-
-- `heatmap` (default),
-- `contour`,
-- `contourf`.
-
-Available plot `type`s for 3D are:
-
-- `contour` (default).
-
-The `alpha` value gets passed to `contour` in 3D.
-"""
-field_plotter(setup; nupdate = 1, kwargs...) =
- processor(state -> field_plot(setup.grid.dimension, setup, state; kwargs...); nupdate)
-
-function field_plot(
- ::Dimension{2},
- setup,
- state;
- fieldname = :vorticity,
- type = heatmap,
- sleeptime = 0.001,
- equal_axis = true,
- displayfig = true,
-)
- (; boundary_conditions, grid) = setup
- (; xlims, ylims, x, y, xp, yp) = grid
-
- if fieldname == :velocity
- xf, yf = xp, yp
- elseif fieldname == :vorticity
- if all(==(:periodic), (boundary_conditions.u.x[1], boundary_conditions.v.y[1]))
- xf = x
- yf = y
- else
- xf = x[2:(end-1)]
- yf = y[2:(end-1)]
- end
- elseif fieldname == :streamfunction
- if boundary_conditions.u.x[1] == :periodic
- xf = x
- else
- xf = x[2:(end-1)]
- end
- if boundary_conditions.v.y[1] == :periodic
- yf = y
- else
- yf = y[2:(end-1)]
- end
- elseif fieldname == :pressure
- error("Not implemented")
- xf, yf = xp, yp
- else
- error("Unknown fieldname")
- end
-
- field = @lift begin
- isnothing(sleeptime) || sleep(sleeptime)
- (; V, p, t) = $state
- f = if fieldname == :velocity
- up, vp = get_velocity(setup, V, t)
- map((u, v) -> √sum(u^2 + v^2), up, vp)
- elseif fieldname == :vorticity
- get_vorticity(setup, V, t)
- elseif fieldname == :streamfunction
- get_streamfunction(setup, V, t)
- elseif fieldname == :pressure
- error("Not implemented")
- reshape(p, length(xp), length(yp))
- end
- Array(f)
- end
-
- lims = @lift begin
- f = $field
- if type == heatmap
- lims = get_lims(f)
- elseif type ∈ (contour, contourf)
- if ≈(extrema(f)...; rtol = 1e-10)
- μ = mean(f)
- a = μ - 1
- b = μ + 1
- f[1] += 1
- f[end] -= 1
- else
- a, b = get_lims(f)
- end
- lims = (a, b)
- end
- lims
- end
-
- fig = Figure()
-
- if type == heatmap
- ax, hm = heatmap(fig[1, 1], xf, yf, field; colorrange = lims)
- elseif type ∈ (contour, contourf)
- ax, hm = type(
- fig[1, 1],
- xf,
- yf,
- field;
- extendlow = :auto,
- extendhigh = :auto,
- levels = @lift(LinRange($(lims)..., 10)),
- colorrange = lims,
- )
- else
- error("Unknown plot type")
- end
-
- ax.title = titlecase(string(fieldname))
- equal_axis && (ax.aspect = DataAspect())
- ax.xlabel = "x"
- ax.ylabel = "y"
- limits!(ax, xlims[1], xlims[2], ylims[1], ylims[2])
- Colorbar(fig[1, 2], hm)
-
- displayfig && display(fig)
-
- fig
-end
-
-function field_plot(
- ::Dimension{3},
- setup,
- state;
- fieldname = :vorticity,
- sleeptime = 0.001,
- alpha = 0.05,
- equal_axis = true,
- levels = 3,
- displayfig = true,
-)
- (; boundary_conditions, grid) = setup
- (; xlims, ylims, x, y, z, xp, yp, zp) = grid
-
- if fieldname == :velocity
- xf, yf, zf = xp, yp, zp
- elseif fieldname == :vorticity
- if all(==(:periodic), (boundary_conditions.u.x[1], boundary_conditions.v.y[1]))
- xf = x
- yf = y
- zf = y
- else
- xf = x[2:(end-1)]
- yf = y[2:(end-1)]
- zf = z[2:(end-1)]
- end
- elseif fieldname == :streamfunction
- if boundary_conditions.u.x[1] == :periodic
- xf = x
- else
- xf = x[2:(end-1)]
- end
- if boundary_conditions.v.y[1] == :periodic
- yf = y
- else
- yf = y[2:(end-1)]
- end
- elseif fieldname == :pressure
- xf, yf, zf = xp, yp, zp
- else
- error("Unknown fieldname")
- end
-
- field = @lift begin
- isnothing(sleeptime) || sleep(sleeptime)
- (; V, p, t) = $state
- f = if fieldname == :velocity
- up, vp, wp = get_velocity(setup, V, t)
- map((u, v, w) -> √sum(u^2 + v^2 + w^2), up, vp, wp)
- elseif fieldname == :vorticity
- get_vorticity(setup, V, t)
- elseif fieldname == :streamfunction
- get_streamfunction(setup, V, t)
- elseif fieldname == :pressure
- reshape(copy(p), length(xp), length(yp), length(zp))
- end
- Array(f)
- end
-
- lims = @lift get_lims($field)
-
- isnothing(levels) && (levels = @lift(LinRange($(lims)..., 10)))
-
- aspect = equal_axis ? (; aspect = :data) : (;)
- fig = Figure()
- ax = Axis3(fig[1, 1]; title = titlecase(string(fieldname)), aspect...)
- hm = contour!(
- ax,
- xf,
- yf,
- zf,
- field;
- levels,
- colorrange = lims,
- shading = false,
- alpha,
- highclip = :red,
- lowclip = :red,
- )
-
- Colorbar(fig[1, 2], hm)
-
- displayfig && display(fig)
-
- fig
-end
-
-"""
- energy_history_plotter(setup)
-
-Create energy history plot, with a history point added every time `step_observer` is updated.
-"""
-energy_history_plotter(setup; nupdate = 1, kwargs...) =
- processor(state -> energy_history_plot(setup, state; kwargs...); nupdate)
-
-function energy_history_plot(setup, state; displayfig = true)
- (; Ωp) = setup.grid
- _points = Point2f[]
- points = @lift begin
- (; V, p, t) = $state
- vels = get_velocity(setup, V, t)
- vels = reshape.(vels, :)
- E = sum(vel -> sum(@. Ωp * vel^2), vels)
- push!(_points, Point2f(t, E))
- end
- fig = lines(points; axis = (; xlabel = "t", ylabel = "Kinetic energy"))
- displayfig && display(fig)
- fig
-end
-
-"""
- energy_spectrum_plotter(setup; nupdate = 1)
-
-Create energy spectrum plot, redrawn every time `step_observer` is updated.
-"""
-energy_spectrum_plotter(setup; nupdate = 1, kwargs...) = processor(
- state -> energy_spectrum_plot(setup.grid.dimension, setup, state; kwargs...);
- nupdate,
-)
-
-function energy_spectrum_plot(::Dimension{2}, setup, state; displayfig = true)
- (; xpp) = setup.grid
- Kx, Ky = size(xpp) .÷ 2
- kx = 1:(Kx-1)
- ky = 1:(Ky-1)
- kk = reshape([sqrt(kx^2 + ky^2) for kx ∈ kx, ky ∈ ky], :)
- ehat = @lift begin
- (; V, p, t) = $state
- up, vp = get_velocity(setup, V, t)
- e = up .^ 2 .+ vp .^ 2
- reshape(abs.(fft(e)[kx.+1, ky.+1]), :)
- end
- espec = Figure()
- ax = Axis(espec[1, 1]; xlabel = "k", ylabel = "e(k)", xscale = log10, yscale = log10)
- ## ylims!(ax, (1e-20, 1))
- scatter!(ax, kk, ehat; label = "Kinetic energy")
- krange = LinRange(extrema(kk)..., 100)
- lines!(ax, krange, 1e7 * krange .^ (-3); label = "k⁻³", color = :red)
- axislegend(ax)
- displayfig && display(espec)
- espec
-end
-
-function energy_spectrum_plot(::Dimension{3}, setup, state, displayfig = true)
- (; xpp) = setup.grid
- Kx, Ky, Kz = size(xpp) .÷ 2
- kx = 1:(Kx-1)
- ky = 1:(Ky-1)
- kz = 1:(Ky-1)
- kk = reshape([sqrt(kx^2 + ky^2 + kz^2) for kx ∈ kx, ky ∈ ky, kz ∈ kz], :)
- ehat = @lift begin
- (; V, p, t) = $state
- V = Array(V)
- up, vp, wp = get_velocity(setup, V, t)
- e = @. up^2 + vp^2 + wp^2
- reshape(abs.(fft(e)[kx.+1, ky.+1, kz.+1]), :)
- end
- espec = Figure()
- ax = Axis(espec[1, 1]; xlabel = "k", ylabel = "e(k)", xscale = log10, yscale = log10)
- ## ylims!(ax, (1e-20, 1))
- scatter!(ax, kk, ehat; label = "Kinetic energy")
- krange = LinRange(extrema(kk)..., 100)
- lines!(ax, krange, 1e6 * krange .^ (-5 / 3); label = "\$k^{-5/3}\$", color = :red)
- axislegend(ax)
- displayfig && display(espec)
- espec
-end
diff --git a/src/setup.jl b/src/setup.jl
index d9c922fa4..542c72950 100644
--- a/src/setup.jl
+++ b/src/setup.jl
@@ -1,142 +1,99 @@
"""
Setup(
- x, y;
- viscosity_model = LaminarModel(; Re = 1000),
- convection_model = NoRegConvectionModel(),
- u_bc = (x, y, t) -> 0,
- v_bc = (x, y, t) -> 0,
- dudt_bc = nothing,
- dvdt_bc = nothing,
- bc_type = (;
- u = (; x = (:periodic, :periodic), y = (:periodic, :periodic)),
- v = (; x = (:periodic, :periodic), y = (:periodic, :periodic)),
- ),
- order4 = false,
- bodyforce_u = (x, y) -> 0,
- bodyforce_v = (x, y) -> 0,
+ x...;
+ boundary_conditions = ntuple(d -> (PeriodicBC(), PeriodicBC()), length(x)),
+ Re = convert(eltype(x[1]), 1_000),
+ bodyforce = nothing,
+ issteadybodyforce = true,
+ closure_model = nothing,
+ ArrayType = Array,
+ workgroupsize = 64,
+ temperature = nothing,
)
-Create 2D setup.
+Create setup.
"""
function Setup(
- x,
- y;
- viscosity_model = LaminarModel(; Re = convert(eltype(x), 1000)),
- convection_model = NoRegConvectionModel(),
- u_bc = (x, y, t) -> 0,
- v_bc = (x, y, t) -> 0,
- dudt_bc = nothing,
- dvdt_bc = nothing,
- bc_type = (;
- u = (; x = (:periodic, :periodic), y = (:periodic, :periodic)),
- v = (; x = (:periodic, :periodic), y = (:periodic, :periodic)),
- ν = (; x = (:dirichlet, :dirichlet), y = (:dirichlet, :dirichlet)),
- ),
- order4 = false,
- bodyforce_u = (x, y) -> 0,
- bodyforce_v = (x, y) -> 0,
- steady_force = true,
+ x...;
+ boundary_conditions = ntuple(d -> (PeriodicBC(), PeriodicBC()), length(x)),
+ Re = convert(eltype(x[1]), 1_000),
+ bodyforce = nothing,
+ issteadybodyforce = true,
+ closure_model = nothing,
+ projectorder = :last,
+ ArrayType = Array,
+ workgroupsize = 64,
+ temperature = nothing,
)
- boundary_conditions =
- BoundaryConditions(u_bc, v_bc; dudt_bc, dvdt_bc, bc_type, T = eltype(x))
- grid = Grid(x, y; boundary_conditions, order4)
- force = SteadyBodyForce(bodyforce_u, bodyforce_v, grid)
- operators = Operators(grid, boundary_conditions, viscosity_model)
- (; grid, boundary_conditions, viscosity_model, convection_model, force, operators)
+ setup = (;
+ grid = Grid(x, boundary_conditions; ArrayType),
+ boundary_conditions,
+ Re,
+ bodyforce,
+ issteadybodyforce = false,
+ closure_model,
+ projectorder,
+ ArrayType,
+ T = eltype(x[1]),
+ workgroupsize,
+ temperature,
+ )
+ if !isnothing(bodyforce) && issteadybodyforce
+ (; dimension, x, N) = setup.grid
+ T = eltype(x[1])
+ F = ntuple(α -> zero(similar(x[1], N)), dimension())
+ u = ntuple(α -> zero(similar(x[1], N)), dimension())
+ bodyforce = bodyforce!(F, u, T(0), setup)
+ setup = (; setup..., issteadybodyforce = true, bodyforce)
+ end
+ setup
end
"""
- Setup(
- x, y, z;
- viscosity_model = LaminarModel(; Re = 1000),
- convection_model = NoRegConvectionModel(),
- u_bc = (x, y, w, t) -> 0.0,
- v_bc = (x, y, w, t) -> 0.0,
- w_bc = (x, y, w, t) -> 0.0,
- dudt_bc = nothing,
- dvdt_bc = nothing,
- dwdt_bc = nothing,
- bc_type = (;
- u = (;
- x = (:periodic, :periodic),
- y = (:periodic, :periodic),
- z = (:periodic, :periodic),
- ),
- v = (;
- x = (:periodic, :periodic),
- y = (:periodic, :periodic),
- z = (:periodic, :periodic),
- ),
- w = (;
- x = (:periodic, :periodic),
- y = (:periodic, :periodic),
- z = (:periodic, :periodic),
- ),
- ν = (;
- x = (:periodic, :periodic),
- y = (:periodic, :periodic),
- z = (:periodic, :periodic),
- ),
- ),
- order4 = false,
- bodyforce_u = (x, y, z) -> 0,
- bodyforce_v = (x, y, z) -> 0,
- bodyforce_w = (x, y, z) -> 0,
+ temperature_equation(;
+ Pr,
+ Ra,
+ Ge,
+ dodissipation = true,
+ boundary_conditions,
+ gdir = 2,
)
-Create 3D setup.
+Temperature equation setup.
+
+The equation is parameterized by three dimensionless numbers (Prandtl number,
+Rayleigh number, and Gebhart number), and requires separate boundary conditions
+for the `temperature` field. The `gdir` keyword specifies the direction gravity,
+while `non_dim_type` specifies the type of non-dimensionalization.
"""
-function Setup(
- x,
- y,
- z;
- viscosity_model = LaminarModel(; Re = convert(eltype(x), 1000)),
- convection_model = NoRegConvectionModel(),
- u_bc = (x, y, w, t) -> 0,
- v_bc = (x, y, w, t) -> 0,
- w_bc = (x, y, w, t) -> 0,
- dudt_bc = nothing,
- dvdt_bc = nothing,
- dwdt_bc = nothing,
- bc_type = (;
- u = (;
- x = (:periodic, :periodic),
- y = (:periodic, :periodic),
- z = (:periodic, :periodic),
- ),
- v = (;
- x = (:periodic, :periodic),
- y = (:periodic, :periodic),
- z = (:periodic, :periodic),
- ),
- w = (;
- x = (:periodic, :periodic),
- y = (:periodic, :periodic),
- z = (:periodic, :periodic),
- ),
- ν = (;
- x = (:periodic, :periodic),
- y = (:periodic, :periodic),
- z = (:periodic, :periodic),
- ),
- ),
- order4 = false,
- bodyforce_u = (x, y, z) -> 0,
- bodyforce_v = (x, y, z) -> 0,
- bodyforce_w = (x, y, z) -> 0,
+function temperature_equation(;
+ Pr,
+ Ra,
+ Ge,
+ dodissipation = true,
+ boundary_conditions,
+ gdir = 2,
+ nondim_type = 1,
)
- boundary_conditions = BoundaryConditions(
- u_bc,
- v_bc,
- w_bc;
- dudt_bc,
- dvdt_bc,
- dwdt_bc,
- bc_type,
- T = eltype(x),
- )
- grid = Grid(x, y, z; boundary_conditions, order4)
- force = SteadyBodyForce(bodyforce_u, bodyforce_v, bodyforce_w, grid)
- operators = Operators(grid, boundary_conditions, viscosity_model)
- (; grid, boundary_conditions, viscosity_model, convection_model, force, operators)
+ if nondim_type == 1
+ # free fall velocity, uref = sqrt(beta*g*Delta T*H)
+ α1 = sqrt(Pr / Ra)
+ α2 = eltype(Pr)(1)
+ α3 = Ge * sqrt(Pr / Ra)
+ α4 = 1 / sqrt(Pr * Ra)
+ elseif nondim_type == 2
+ # uref = kappa/H (based on heat conduction time scale)
+ α1 = Pr
+ α2 = Pr * Ra
+ α3 = Ge / Ra
+ α4 = 1
+ elseif nondim_type == 3
+ # uref = sqrt(c*Delta T)
+ α1 = sqrt(Pr * Ge / Ra)
+ α2 = Ge
+ α3 = sqrt(Pr * Ge / Ra)
+ α4 = sqrt(Ge / (Pr * Ra))
+ end
+ γ = α1 / α3
+ (; α1, α2, α3, α4, γ, dodissipation, boundary_conditions, gdir)
end
diff --git a/src/solvers/get_timestep.jl b/src/solvers/get_timestep.jl
index 6097db59a..c17ae5ef6 100644
--- a/src/solvers/get_timestep.jl
+++ b/src/solvers/get_timestep.jl
@@ -24,33 +24,33 @@ function get_timestep(::Dimension{2}, stepper, cfl)
vₕ = @view V[indv]
# For explicit methods only
- if isexplicit(method)
- # Convective part
- Cu =
- Cux * spdiagm(Iu_ux * uₕ + yIu_ux) * Au_ux +
- Cuy * spdiagm(Iv_uy * vₕ + yIv_uy) * Au_uy
- Cv =
- Cvx * spdiagm(Iu_vx * uₕ + yIu_vx) * Av_vx +
- Cvy * spdiagm(Iv_vy * vₕ + yIv_vy) * Av_vy
- C = blockdiag(Cu, Cv)
- test = spdiagm(1 ./ Ω) * C
- sum_conv = abs.(test) * ones(T, NV) - diag(abs.(test)) - diag(test)
- λ_conv = maximum(sum_conv)
-
- # Based on max. value of stability region (not a very good indication
- # For the methods that do not include the imaginary axis)
- Δt_conv = lambda_conv_max(method) / λ_conv
-
- ## Diffusive part
- test = Diagonal(1 ./ Ω) * Diff
- sum_diff = abs.(test) * ones(T, NV) - diag(abs.(test)) - diag(test)
- λ_diff = maximum(sum_diff)
-
- # Based on max. value of stability region
- Δt_diff = lambda_diff_max(method) / λ_diff
-
- Δt = cfl * min(Δt_conv, Δt_diff)
- end
+ @assert isexplicit(method) "Adaptive timestep requires explicit method"
+
+ # Convective part
+ Cu =
+ Cux * spdiagm(Iu_ux * uₕ + yIu_ux) * Au_ux +
+ Cuy * spdiagm(Iv_uy * vₕ + yIv_uy) * Au_uy
+ Cv =
+ Cvx * spdiagm(Iu_vx * uₕ + yIu_vx) * Av_vx +
+ Cvy * spdiagm(Iv_vy * vₕ + yIv_vy) * Av_vy
+ C = blockdiag(Cu, Cv)
+ test = spdiagm(1 ./ Ω) * C
+ sum_conv = abs.(test) * ones(T, NV) - diag(abs.(test)) - diag(test)
+ λ_conv = maximum(sum_conv)
+
+ # Based on max. value of stability region (not a very good indication
+ # For the methods that do not include the imaginary axis)
+ Δt_conv = lambda_conv_max(method) / λ_conv
+
+ ## Diffusive part
+ test = Diagonal(1 ./ Ω) * Diff
+ sum_diff = abs.(test) * ones(T, NV) - diag(abs.(test)) - diag(test)
+ λ_diff = maximum(sum_diff)
+
+ # Based on max. value of stability region
+ Δt_diff = lambda_diff_max(method) / λ_diff
+
+ Δt = cfl * min(Δt_conv, Δt_diff)
Δt
end
@@ -75,39 +75,39 @@ function get_timestep(::Dimension{3}, stepper, cfl)
wₕ = @view V[indw]
# For explicit methods only
- if isexplicit(method)
- # Convective part
- Cu =
- Cux * spdiagm(Iu_ux * uₕ + yIu_ux) * Au_ux +
- Cuy * spdiagm(Iv_uy * vₕ + yIv_uy) * Au_uy +
- Cuz * spdiagm(Iw_uz * wₕ + yIw_uz) * Au_uz
- Cv =
- Cvx * spdiagm(Iu_vx * uₕ + yIu_vx) * Av_vx +
- Cvy * spdiagm(Iv_vy * vₕ + yIv_vy) * Av_vy +
- Cvz * spdiagm(Iw_vz * wₕ + yIw_vz) * Av_vz
- Cw =
- Cwx * spdiagm(Iu_wx * uₕ + yIu_wx) * Aw_wx +
- Cwy * spdiagm(Iv_wy * vₕ + yIv_wy) * Aw_wy +
- Cwz * spdiagm(Iw_wz * wₕ + yIw_wz) * Aw_wz
- C = blockdiag(Cu, Cv, Cw)
- test = spdiagm(1 ./ Ω) * C
- sum_conv = abs.(test) * ones(T, NV) - diag(abs.(test)) - diag(test)
- λ_conv = maximum(sum_conv)
-
- # Based on max. value of stability region (not a very good indication
- # For the methods that do not include the imaginary axis)
- Δt_conv = lambda_conv_max(method) / λ_conv
-
- ## Diffusive part
- test = Diagonal(1 ./ Ω) * Diff
- sum_diff = abs.(test) * ones(T, NV) - diag(abs.(test)) - diag(test)
- λ_diff = maximum(sum_diff)
-
- # Based on max. value of stability region
- Δt_diff = lambda_diff_max(method) / λ_diff
-
- Δt = cfl * min(Δt_conv, Δt_diff)
- end
+ @assert isexplicit(method) "Adaptive timestep requires explicit method"
+
+ # Convective part
+ Cu =
+ Cux * spdiagm(Iu_ux * uₕ + yIu_ux) * Au_ux +
+ Cuy * spdiagm(Iv_uy * vₕ + yIv_uy) * Au_uy +
+ Cuz * spdiagm(Iw_uz * wₕ + yIw_uz) * Au_uz
+ Cv =
+ Cvx * spdiagm(Iu_vx * uₕ + yIu_vx) * Av_vx +
+ Cvy * spdiagm(Iv_vy * vₕ + yIv_vy) * Av_vy +
+ Cvz * spdiagm(Iw_vz * wₕ + yIw_vz) * Av_vz
+ Cw =
+ Cwx * spdiagm(Iu_wx * uₕ + yIu_wx) * Aw_wx +
+ Cwy * spdiagm(Iv_wy * vₕ + yIv_wy) * Aw_wy +
+ Cwz * spdiagm(Iw_wz * wₕ + yIw_wz) * Aw_wz
+ C = blockdiag(Cu, Cv, Cw)
+ test = spdiagm(1 ./ Ω) * C
+ sum_conv = abs.(test) * ones(T, NV) - diag(abs.(test)) - diag(test)
+ λ_conv = maximum(sum_conv)
+
+ # Based on max. value of stability region (not a very good indication
+ # For the methods that do not include the imaginary axis)
+ Δt_conv = lambda_conv_max(method) / λ_conv
+
+ ## Diffusive part
+ test = Diagonal(1 ./ Ω) * Diff
+ sum_diff = abs.(test) * ones(T, NV) - diag(abs.(test)) - diag(test)
+ λ_diff = maximum(sum_diff)
+
+ # Based on max. value of stability region
+ Δt_diff = lambda_diff_max(method) / λ_diff
+
+ Δt = cfl * min(Δt_conv, Δt_diff)
Δt
end
diff --git a/src/solvers/pressure/pressure_additional_solve.jl b/src/solvers/pressure/pressure_additional_solve.jl
deleted file mode 100644
index f55c6b330..000000000
--- a/src/solvers/pressure/pressure_additional_solve.jl
+++ /dev/null
@@ -1,71 +0,0 @@
-"""
- pressure_additional_solve(pressure_solver, V, p, t, setup; bc_vectors = nothing)
-
-Do additional pressure solve. This makes the pressure compatible with the velocity
-field, resulting in same order pressure as velocity.
-"""
-function pressure_additional_solve(pressure_solver, V, p, t, setup; bc_vectors = nothing)
- (; grid, operators, boundary_conditions) = setup
- (; bc_unsteady) = boundary_conditions
- (; Ω) = grid
- (; M) = operators
-
- # Get updated BC for ydM (time derivative of BC in ydM)
- # FIXME: `get_bc_vectors` are called to often (also inside `momentum!`)
- if isnothing(bc_vectors) || bc_unsteady
- bc_vectors = get_bc_vectors(setup, t)
- end
- (; ydM) = bc_vectors
-
- # Momentum already contains G*p with the current p, we therefore effectively solve for
- # the pressure difference
- F, = momentum(V, V, p, t, setup; bc_vectors)
- f = M * (1 ./ Ω .* F) + ydM
-
- Δp = pressure_poisson(pressure_solver, f)
- p .+ Δp
-end
-
-"""
- pressure_additional_solve!(pressure_solver, V, p, t, setup, momentum_cache, F, f; bc_vectors)
-
-Do additional pressure solve. This makes the pressure compatible with the velocity
-field, resulting in same order pressure as velocity.
-"""
-function pressure_additional_solve!(
- pressure_solver,
- V,
- p,
- t,
- setup,
- momentum_cache,
- F,
- f,
- Δp;
- bc_vectors,
-)
- (; grid, operators, boundary_conditions) = setup
- (; bc_unsteady) = boundary_conditions
- (; Ω) = grid
- (; M) = operators
-
- # Get updated BC for ydM (time derivative of BC in ydM)
- # FIXME: `get_bc_vectors` are called to often (also inside `momentum!`)
- if isnothing(bc_vectors) || bc_unsteady
- bc_vectors = get_bc_vectors(setup, t)
- end
- (; ydM) = bc_vectors
-
- # Momentum already contains G*p with the current p, we therefore effectively solve for
- # the pressure difference
- momentum!(F, nothing, V, V, p, t, setup, momentum_cache; bc_vectors)
-
- # f = M * (1 ./ Ω .* F) + ydM
- @. F = 1 ./ Ω .* F
- mul!(f, M, F)
- @. f = f + ydM
-
- pressure_poisson!(pressure_solver, Δp, f)
-
- p .+= Δp
-end
diff --git a/src/solvers/pressure/pressure_poisson.jl b/src/solvers/pressure/pressure_poisson.jl
deleted file mode 100644
index 81c82db29..000000000
--- a/src/solvers/pressure/pressure_poisson.jl
+++ /dev/null
@@ -1,83 +0,0 @@
-"""
- pressure_poisson(solver, f)
-
-Solve the Poisson equation for the pressure with right hand side `f` at time `t`.
-For periodic and no-slip BC, the sum of `f` should be zero.
-
-Non-mutating/allocating/out-of-place version.
-
-See also [`pressure_poisson!`](@ref).
-"""
-function pressure_poisson end
-
-function pressure_poisson(solver::DirectPressureSolver, f)
- # Assume the Laplace matrix is known (A) and is possibly factorized
-
- # Use pre-determined decomposition
- solver.A_fact \ f
-end
-
-function pressure_poisson(solver::CGPressureSolver, f)
- (; A, abstol, reltol, maxiter) = solver
- cg(A, f; abstol, reltol, maxiter)
-end
-
-function pressure_poisson(solver::SpectralPressureSolver, f)
- (; Ahat) = solver
-
- f = reshape(f, size(Ahat))
-
- # Fourier transform of right hand side
- fhat = fft(f)
-
- # Solve for coefficients in Fourier space
- phat = @. -fhat / Ahat
-
- # Transform back
- p = ifft(phat)
-
- reshape(real.(p), :)
-end
-
-"""
- pressure_poisson!(solver, p, f)
-
-Solve the Poisson equation for the pressure with right hand side `f` at time `t`.
-For periodic and no-slip BC, the sum of `f` should be zero.
-
-Mutating/non-allocating/in-place version.
-
-See also [`pressure_poisson`](@ref).
-"""
-function pressure_poisson! end
-
-function pressure_poisson!(solver::DirectPressureSolver, p, f)
- # Assume the Laplace matrix is known (A) and is possibly factorized
-
- # Use pre-determined decomposition
- p .= solver.A_fact \ f
-end
-
-function pressure_poisson!(solver::CGPressureSolver, p, f)
- # TODO: Preconditioner
- (; A, abstol, reltol, maxiter) = solver
- cg!(p, A, f; abstol, reltol, maxiter)
-end
-
-function pressure_poisson!(solver::SpectralPressureSolver, p, f)
- (; Ahat, fhat, phat) = solver
-
- fhat[:] .= complex.(f)
-
- # Fourier transform of right hand side
- fft!(fhat)
-
- # Solve for coefficients in Fourier space
- @. phat = -fhat / Ahat
-
- # Transform back
- ifft!(phat)
- @. p[:] = real(@view phat[:])
-
- p
-end
diff --git a/src/solvers/pressure/pressure_solvers.jl b/src/solvers/pressure/pressure_solvers.jl
deleted file mode 100644
index 33f864e6f..000000000
--- a/src/solvers/pressure/pressure_solvers.jl
+++ /dev/null
@@ -1,164 +0,0 @@
-"""
- AbstractPressureSolver
-
-Pressure solver for the Poisson equation.
-"""
-abstract type AbstractPressureSolver{T} end
-
-"""
- DirectPressureSolver()
-
-Direct pressure solver using a LU decomposition.
-"""
-struct DirectPressureSolver{T,F<:Factorization{T}} <: AbstractPressureSolver{T}
- A_fact::F
- function DirectPressureSolver(setup)
- (; A) = setup.operators
- T = eltype(A)
- fact = factorize(setup.operators.A)
- new{T,typeof(fact)}(fact)
- end
-end
-
-# This moves all the inner arrays to the GPU when calling
-# `cu(::SpectralPressureSolver)` from CUDA.jl
-# TODO: CUDA does not support `factorize`, `lu`, etc, but maybe `L` and `U` can be
-# converted to `CuArray` after factorization on the CPU
-Adapt.adapt_structure(to, s::DirectPressureSolver) = error(
- "`DirectPressureSolver` is not yet implemented for CUDA. Consider using `CGPressureSolver`.",
-)
-
-"""
- CGPressureSolver(setup; [abstol], [reltol], [maxiter])
-
-Conjugate gradients iterative pressure solver.
-"""
-struct CGPressureSolver{T,M<:AbstractMatrix{T}} <: AbstractPressureSolver{T}
- A::M
- abstol::T
- reltol::T
- maxiter::Int
-end
-
-function CGPressureSolver(
- setup;
- abstol = 0,
- reltol = √eps(eltype(setup.operators.A)),
- maxiter = size(setup.operators.A, 2),
-)
- (; A) = setup.operators
- CGPressureSolver{eltype(A),typeof(A)}(A, abstol, reltol, maxiter)
-end
-
-# This moves all the inner arrays to the GPU when calling
-# `cu(::SpectralPressureSolver)` from CUDA.jl
-Adapt.adapt_structure(to, s::CGPressureSolver) = CGPressureSolver(
- adapt(to, s.A),
- adapt(to, s.abstol),
- adapt(to, s.reltol),
- adapt(to, s.maxiter),
-)
-
-struct SpectralPressureSolver{T,A<:AbstractArray{Complex{T}}} <: AbstractPressureSolver{T}
- Ahat::A
- phat::A
- fhat::A
-end
-
-# This moves all the inner arrays to the GPU when calling
-# `cu(::SpectralPressureSolver)` from CUDA.jl
-Adapt.adapt_structure(to, s::SpectralPressureSolver) =
- SpectralPressureSolver(adapt(to, s.Ahat), adapt(to, s.phat), adapt(to, s.fhat))
-
-"""
- SpectralPressureSolver(setup)
-
-Build spectral pressure solver from setup.
-"""
-SpectralPressureSolver(setup) = SpectralPressureSolver(setup.grid.dimension, setup)
-
-function SpectralPressureSolver(::Dimension{2}, setup)
- (; grid, boundary_conditions) = setup
- (; hx, hy, Npx, Npy) = grid
-
- T = eltype(hx)
- AT = typeof(hx)
-
- if any(
- !isequal((:periodic, :periodic)),
- (boundary_conditions.u.x, boundary_conditions.v.y),
- )
- error("SpectralPressureSolver only implemented for periodic boundary conditions")
- end
-
- Δx = first(hx)
- Δy = first(hy)
- if any(!≈(Δx), hx) || any(!≈(Δy), hy)
- error("SpectralPressureSolver requires uniform grid along each dimension")
- end
-
- # Fourier transform of the discretization
- # Assuming uniform grid, although Δx and Δy do not need to be the same
- i = AT(0:(Npx-1))
- j = reshape(AT(0:(Npy-1)), 1, :)
-
- # Scale with Δx*Δy, since we solve the PDE in integrated form
- Ahat = @. 4 * Δx * Δy * (sin(i * π / Npx)^2 / Δx^2 + sin(j * π / Npy)^2 / Δy^2)
-
- # Pressure is determined up to constant, fix at 0
- Ahat[1] = 1
-
- Ahat = complex(Ahat)
-
- # Placeholders for intermediate results
- phat = similar(Ahat)
- fhat = similar(Ahat)
-
- SpectralPressureSolver{T,typeof(Ahat)}(Ahat, phat, fhat)
-end
-
-function SpectralPressureSolver(::Dimension{3}, setup)
- (; grid, boundary_conditions) = setup
- (; hx, hy, hz, Npx, Npy, Npz) = grid
-
- T = eltype(hx)
- AT = typeof(hx)
-
- if any(
- !isequal((:periodic, :periodic)),
- [boundary_conditions.u.x, boundary_conditions.v.y, boundary_conditions.w.z],
- )
- error("SpectralPressureSolver only implemented for periodic boundary conditions")
- end
-
- Δx = first(hx)
- Δy = first(hy)
- Δz = first(hz)
- if any(!≈(Δx), hx) || any(!≈(Δy), hy) || any(!≈(Δz), hz)
- error("SpectralPressureSolver requires uniform grid along each dimension")
- end
-
- # Fourier transform of the discretization
- # Assuming uniform grid, although Δx and Δy do not need to be the same
- i = AT(0:(Npx-1))
- j = reshape(AT(0:(Npy-1)), 1, :)
- k = reshape(AT(0:(Npz-1)), 1, 1, :)
-
- # Scale with Δx*Δy*Δz, since we solve the PDE in integrated form
- Ahat = @. 4 *
- Δx *
- Δy *
- Δz *
- (sin(i * π / Npx)^2 / Δx^2 + sin(j * π / Npy)^2 / Δy^2 + sin(k * π / Npz)^2 / Δz^2)
-
- # Pressure is determined up to constant, fix at 0
- Ahat[1] = 1
-
- Ahat = complex(Ahat)
-
- # Placeholders for intermediate results
- phat = similar(Ahat)
- fhat = similar(Ahat)
-
- SpectralPressureSolver{T,typeof(Ahat)}(Ahat, phat, fhat)
-end
diff --git a/src/solvers/solve_unsteady.jl b/src/solvers/solve_unsteady.jl
index 00ff3e75a..a95445cdf 100644
--- a/src/solvers/solve_unsteady.jl
+++ b/src/solvers/solve_unsteady.jl
@@ -1,14 +1,16 @@
"""
- function solve_unsteady(
- setup, V₀, p₀, tlims;
- method = RK44(; T = eltype(V₀)),
- pressure_solver = DirectPressureSolver(setup),
- Δt = nothing,
+ solve_unsteady(;
+ setup,
+ tlims,
+ ustart,
+ tempstart = nothing,
+ method = RKMethods.RK44(; T = eltype(u₀[1])),
+ psolver = default_psolver(setup),
+ Δt = zero(eltype(u₀[1])),
cfl = 1,
n_adapt_Δt = 1,
- inplace = false,
- processors = (),
- device = identity,
+ docopy = true,
+ processors = (;),
)
Solve unsteady problem using `method`.
@@ -18,123 +20,87 @@ an integer.
If `Δt = nothing`, the time step is chosen every `n_adapt_Δt` iteration with
CFL-number `cfl` .
-Each `processor` is called after every `processor.nupdate` time step.
-
-All arrays and operators are passed through the `device` function.
-This allows for performing computations on a different device than the host (CPU).
-To compute on an Nvidia GPU using CUDA, change
-
-```
-solve_unsteady(setup, V₀, p₀, tlims; kwargs...)
-```
-
-to the following:
-
-```
-using CUDA
-solve_unsteady(
- setup, V₀, p₀, tlims;
- device = cu,
- kwargs...
-)
-```
+The `processors` are called after every time step.
Note that the `state` observable passed to the `processor.initialize` function
contains vector living on the device, and you may have to move them back to
-the host using `Array(V)` and `Array(p)` in the processor.
+the host using `Array(u)` in the processor.
+
+Return `(; u, t), outputs`, where `outputs` is a named tuple with the
+outputs of `processors` with the same field names.
"""
-function solve_unsteady(
+function solve_unsteady(;
setup,
- V₀,
- p₀,
- tlims;
- method = RK44(; T = eltype(V₀)),
- pressure_solver = DirectPressureSolver(setup),
- Δt = zero(eltype(V₀)),
+ tlims,
+ ustart,
+ tempstart = nothing,
+ method = RKMethods.RK44(; T = eltype(ustart[1])),
+ psolver = default_psolver(setup),
+ Δt = zero(eltype(ustart[1])),
cfl = 1,
n_adapt_Δt = 1,
- inplace = false,
- processors = (),
- device = identity,
+ docopy = true,
+ processors = (;),
+ θ = nothing,
)
- t_start, t_end = tlims
- isadaptive = isnothing(Δt)
- if !isadaptive
- nstep = round(Int, (t_end - t_start) / Δt)
- Δt = (t_end - t_start) / nstep
- end
+ docopy && (ustart = copy.(ustart))
+ docopy && !isnothing(tempstart) && (tempstart = copy(tempstart))
- # Initialize BC arrays (currently only done on host, due to Kronecker
- # products)
- bc_vectors = get_bc_vectors(setup, t_start)
+ tstart, tend = tlims
+ isadaptive = isnothing(Δt)
- # Move vectors and operators to device (if any).
- setup = device(setup)
- V₀ = device(V₀)
- p₀ = device(p₀)
- bc_vectors = device(bc_vectors)
- pressure_solver = device(pressure_solver)
-
- if inplace
- cache = ode_method_cache(method, setup, V₀, p₀)
- momentum_cache = MomentumCache(setup, V₀, p₀)
- end
+ # Cache arrays for intermediate computations
+ cache = ode_method_cache(method, setup, ustart, tempstart)
# Time stepper
- stepper = create_stepper(
- method;
- setup,
- pressure_solver,
- bc_vectors,
- V = copy(V₀),
- p = copy(p₀),
- t = t_start,
- )
-
- # Get initial time step
- isadaptive && (Δt = get_timestep(stepper, cfl))
+ stepper =
+ create_stepper(method; setup, psolver, u = ustart, temp = tempstart, t = tstart)
# Initialize processors for iteration results
- state = get_state(stepper)
- states = map(ps -> Observable(state), processors)
- initialized = map((ps, o) -> ps.initialize(o), processors, states)
+ state = Observable(get_state(stepper))
+ initialized = (; (k => v.initialize(state) for (k, v) in pairs(processors))...)
- while stepper.t < t_end
- if isadaptive
+ if isadaptive
+ while stepper.t < tend
if stepper.n % n_adapt_Δt == 0
# Change timestep based on operators
Δt = get_timestep(stepper, cfl)
end
# Make sure not to step past `t_end`
- Δt = min(Δt, t_end - stepper.t)
- end
+ Δt = min(Δt, tend - stepper.t)
- # Perform a single time step with the time integration method
- if inplace
- stepper = step!(method, stepper, Δt; cache, momentum_cache)
- else
- stepper = step(method, stepper, Δt)
- end
+ # Perform a single time step with the time integration method
+ stepper = timestep!(method, stepper, Δt; θ, cache)
- # Process iteration results with each processor
- for (ps, o) ∈ zip(processors, states)
- # Only update each `nupdate`-th iteration
- stepper.n % ps.nupdate == 0 && (o[] = get_state(stepper))
+ # Process iteration results with each processor
+ state[] = get_state(stepper)
+ end
+ else
+ nstep = round(Int, (tend - tstart) / Δt)
+ Δt = (tend - tstart) / nstep
+ for it = 1:nstep
+ # Perform a single time step with the time integration method
+ stepper = timestep!(method, stepper, Δt; θ, cache)
+
+ # Process iteration results with each processor
+ state[] = get_state(stepper)
end
end
- (; V, p, t, n) = stepper
- finalized = map((ps, i) -> ps.finalize(i, get_state(stepper)), processors, initialized)
-
# Final state
- (; V, p) = stepper
+ (; u, temp, t) = stepper
+
+ # Processor outputs
+ outputs = (;
+ (k => processors[k].finalize(initialized[k], state) for k in keys(processors))...
+ )
- # Move output arrays to host
- Array(V), Array(p), finalized
+ # Return state and outputs
+ (; u, temp, t), outputs
end
function get_state(stepper)
- (; V, p, t, n) = stepper
- (; V, p, t, n)
+ (; u, temp, t, n) = stepper
+ (; u, temp, t, n)
end
diff --git a/src/time_steppers/methods.jl b/src/time_steppers/methods.jl
index 15b5eeb83..7921872cb 100644
--- a/src/time_steppers/methods.jl
+++ b/src/time_steppers/methods.jl
@@ -12,44 +12,18 @@ abstract type AbstractODEMethod{T} end
α₂ = T(-1 // 2),
θ = T(1 // 2),
p_add_solve = true,
- method_startup = RK44(; T),
+ method_startup,
)
-IMEX AB-CN: Adams-Bashforth for explicit convection (parameters `α₁` and `α₂`) and
-Crank-Nicolson for implicit diffusion (implicitness `θ`).
-The method is second order for `θ = 1/2`.
-
-Adams-Bashforth for convection and Crank-Nicolson for diffusion
-formulation:
-
-```math
-\\begin{align*}
-(\\mathbf{u}^{n+1} - \\mathbf{u}^n) / Δt & =
- -(\\alpha_1 \\mathbf{c}^n + \\alpha_2 \\mathbf{c}^{n-1}) \\\\
- & + \\theta \\mathbf{d}^{n+1} + (1-\\theta) \\mathbf{d}^n \\\\
- & + \\theta \\mathbf{F}^{n+1} + (1-\\theta) \\mathbf{F}^n \\\\
- & + \\theta \\mathbf{BC}^{n+1} + (1-\\theta) \\mathbf{BC}^n \\\\
- & - \\mathbf{G} \\mathbf{p} + \\mathbf{y}_p
-\\end{align*}
-```
-
-where BC are boundary conditions of diffusion. This is rewritten as:
-
-```math
-\\begin{align*}
-(\\frac{1}{\\Delta t} \\mathbf{I} - \\theta \\mathbf{D}) \\mathbf{u}^{n+1} & =
- (\\frac{1}{\\Delta t} \\mathbf{I} - (1 - \\theta) \\mathbf{D}) \\mathbf{u}^{n} \\\\
- & - (\\alpha_1 \\mathbf{c}^n + \\alpha_2 \\mathbf{c}^{n-1}) \\\\
- & + \\theta \\mathbf{F}^{n+1} + (1-\\theta) \\mathbf{F}^{n} \\\\
- & + \\theta \\mathbf{BC}^{n+1} + (1-\\theta) \\mathbf{BC}^{n} \\\\
- & - \\mathbf{G} \\mathbf{p} + \\mathbf{y}_p
-\\end{align*}
-```
-
-The LU decomposition of the LHS matrix is computed every time the time step changes.
-
-Note that, in constrast to explicit methods, the pressure from previous time steps has an
-influence on the accuracy of the velocity.
+IMEX AB-CN: Adams-Bashforth for explicit convection (parameters `α₁` and `α₂`)
+and Crank-Nicolson for implicit diffusion (implicitness `θ`). The method is
+second order for `θ = 1/2`.
+
+The LU decomposition of the LHS matrix is computed every time the time step
+changes.
+
+Note that, in contrast to explicit methods, the pressure from previous time
+steps has an influence on the accuracy of the velocity.
"""
struct AdamsBashforthCrankNicolsonMethod{T,M} <: AbstractODEMethod{T}
α₁::T
@@ -63,7 +37,7 @@ struct AdamsBashforthCrankNicolsonMethod{T,M} <: AbstractODEMethod{T}
α₂ = T(-1 // 2),
θ = T(1 // 2),
p_add_solve = true,
- method_startup = RK44(; T),
+ method_startup,
) = new{T,typeof(method_startup)}(α₁, α₂, θ, p_add_solve, method_startup)
end
@@ -72,30 +46,19 @@ end
T = Float64;
β = T(1 // 2),
p_add_solve = true,
- method_startup = RK44(; T),
+ method_startup,
)
Explicit one-leg β-method following symmetry-preserving discretization of
-turbulent flow. See [Verstappen and Veldman (JCP 2003)] for details, or [Direct numerical
-simulation of turbulence at lower costs (Journal of Engineering Mathematics 1997)].
-
-Formulation:
-
-```math
-\\frac{(\\beta + 1/2) u^{n+1} - 2 \\beta u^{n} + (\\beta - 1/2) u^{n-1}}{\\Delta t} = F((1 +
-\\beta) u^n - \\beta u^{n-1}).
-```
+turbulent flow. See Verstappen and Veldman [Verstappen2003](@cite)
+[Verstappen1997](@cite) for details.
"""
struct OneLegMethod{T,M} <: AbstractODEMethod{T}
β::T
p_add_solve::Bool
method_startup::M
- OneLegMethod(
- T = Float64;
- β = T(1 // 2),
- p_add_solve = true,
- method_startup = RK44(; T),
- ) = new{T,typeof(method_startup)}(β, p_add_solve, method_startup)
+ OneLegMethod(T = Float64; β = T(1 // 2), p_add_solve = true, method_startup) =
+ new{T,typeof(method_startup)}(β, p_add_solve, method_startup)
end
"""
diff --git a/src/time_steppers/step.jl b/src/time_steppers/step.jl
index fbc99169d..de42789d5 100644
--- a/src/time_steppers/step.jl
+++ b/src/time_steppers/step.jl
@@ -5,9 +5,9 @@ Perform one time step.
Non-mutating/allocating/out-of-place version.
-See also [`step!`](@ref).
+See also [`timestep!`](@ref).
"""
-function step end
+function timestep end
# step(stepper, Δt; bc_vectors = nothing) = step(stepper.method, stepper, Δt; bc_vectors = nothing)
"""
@@ -17,9 +17,9 @@ Perform one time step>
Mutating/non-allocating/in-place version.
-See also [`step`](@ref).
+See also [`timestep`](@ref).
"""
-function step! end
+function timestep! end
# step!(stepper, Δt; kwargs...) = step!(stepper.method, stepper, Δt; kwargs...)
include("step_ab_cn.jl")
diff --git a/src/time_steppers/step_ab_cn.jl b/src/time_steppers/step_ab_cn.jl
index 2266257ae..c9df17c48 100644
--- a/src/time_steppers/step_ab_cn.jl
+++ b/src/time_steppers/step_ab_cn.jl
@@ -1,7 +1,7 @@
create_stepper(
method::AdamsBashforthCrankNicolsonMethod;
setup,
- pressure_solver,
+ psolver,
bc_vectors,
V,
p,
@@ -14,18 +14,17 @@ create_stepper(
cₙ = copy(V),
tₙ = t,
Diff_fact = spzeros(eltype(V), 0, 0),
-) = (; setup, pressure_solver, bc_vectors, V, p, t, n, Vₙ, pₙ, cₙ, tₙ, Diff_fact)
+) = (; setup, psolver, bc_vectors, V, p, t, n, Vₙ, pₙ, cₙ, tₙ, Diff_fact)
-function step(method::AdamsBashforthCrankNicolsonMethod, stepper, Δt)
- (; setup, pressure_solver, bc_vectors, V, p, t, n, Vₙ, pₙ, cₙ, tₙ, Diff_fact) = stepper
- (; convection_model, viscosity_model, force, grid, operators, boundary_conditions) =
+function timestep(method::AdamsBashforthCrankNicolsonMethod, stepper, Δt)
+ (; setup, psolver, bc_vectors, V, p, t, n, Vₙ, pₙ, cₙ, tₙ, Diff_fact) = stepper
+ (; convection_model, viscosity_model, Re, force, grid, operators, boundary_conditions) =
setup
(; bc_unsteady) = boundary_conditions
(; NV, Ω) = grid
(; G, M) = operators
(; Diff) = operators
(; p_add_solve, α₁, α₂, θ, method_startup) = method
- (; Re) = viscosity_model
T = typeof(Δt)
@@ -33,7 +32,7 @@ function step(method::AdamsBashforthCrankNicolsonMethod, stepper, Δt)
# the first iteration. Do one startup step instead
if n == 0
stepper_startup =
- create_stepper(method_startup; setup, pressure_solver, bc_vectors, V, p, t)
+ create_stepper(method_startup; setup, psolver, bc_vectors, V, p, t)
n += 1
Vₙ = V
pₙ = p
@@ -48,11 +47,11 @@ function step(method::AdamsBashforthCrankNicolsonMethod, stepper, Δt)
Diff_fact = lu(I(NV) - θ * Δt / Re * Diagonal(1 ./ Ω) * Diff)
end
- (; V, p, t) = step(method_startup, stepper_startup, Δt)
+ (; V, p, t) = timestep(method_startup, stepper_startup, Δt)
return create_stepper(
method;
setup,
- pressure_solver,
+ psolver,
bc_vectors,
V,
p,
@@ -130,7 +129,7 @@ function step(method::AdamsBashforthCrankNicolsonMethod, stepper, Δt)
f = (M * V + yM) / Δt - M * y_Δp
# Solve the Poisson equation for the pressure
- Δp = pressure_poisson(pressure_solver, f)
+ Δp = poisson(psolver, f)
# Update velocity field
V -= Δt ./ Ω .* (G * Δp .+ y_Δp)
@@ -139,7 +138,7 @@ function step(method::AdamsBashforthCrankNicolsonMethod, stepper, Δt)
p = pₙ .+ Δp
if p_add_solve
- p = pressure_additional_solve(pressure_solver, V, p, tₙ + Δt, setup; bc_vectors)
+ p = pressure(psolver, V, p, tₙ + Δt, setup; bc_vectors)
end
t = tₙ + Δtₙ
@@ -147,7 +146,7 @@ function step(method::AdamsBashforthCrankNicolsonMethod, stepper, Δt)
create_stepper(
method;
setup,
- pressure_solver,
+ psolver,
bc_vectors,
V,
p,
@@ -161,15 +160,15 @@ function step(method::AdamsBashforthCrankNicolsonMethod, stepper, Δt)
)
end
-function step!(
+function timestep!(
method::AdamsBashforthCrankNicolsonMethod,
stepper,
Δt;
cache,
momentum_cache,
)
- (; setup, pressure_solver, bc_vectors, V, p, t, n, Vₙ, pₙ, cₙ, tₙ, Diff_fact) = stepper
- (; convection_model, viscosity_model, force, grid, operators, boundary_conditions) =
+ (; setup, psolver, bc_vectors, V, p, t, n, Vₙ, pₙ, cₙ, tₙ, Diff_fact) = stepper
+ (; convection_model, viscosity_model, Re, force, grid, operators, boundary_conditions) =
setup
(; bc_unsteady) = boundary_conditions
(; NV, Ω) = grid
@@ -178,7 +177,6 @@ function step!(
(; p_add_solve, α₁, α₂, θ, method_startup) = method
(; cₙ₋₁, F, f, Δp, Rr, b, bₙ, bₙ₊₁, yDiffₙ, yDiffₙ₊₁, Gpₙ) = cache
(; d, ∇d) = momentum_cache
- (; Re) = viscosity_model
T = typeof(Δt)
@@ -186,7 +184,7 @@ function step!(
# the first iteration. Do one startup step instead
if n == 0
stepper_startup =
- create_stepper(method_startup; setup, pressure_solver, bc_vectors, V, p, t)
+ create_stepper(method_startup; setup, psolver, bc_vectors, V, p, t)
n += 1
Vₙ = V
pₙ = p
@@ -200,11 +198,11 @@ function step!(
Diff_fact = lu(I(NV) - θ * Δt / Re * Diagonal(1 ./ Ω) * Diff)
# Note: We do one out-of-place step here, with a few allocations
- (; V, p, t) = step(method_startup, stepper_startup, Δt)
+ (; V, p, t) = timestep(method_startup, stepper_startup, Δt)
return create_stepper(
method;
setup,
- pressure_solver,
+ psolver,
bc_vectors,
V,
p,
@@ -285,7 +283,7 @@ function step!(
f = (M * V + yM) / Δt - M * y_Δp
# Solve the Poisson equation for the pressure
- pressure_poisson!(pressure_solver, Δp, f)
+ poisson!(psolver, Δp, f)
# Update velocity field
V .-= Δt ./ Ω .* (G * Δp .+ y_Δp)
@@ -294,18 +292,7 @@ function step!(
p .= pₙ .+ Δp
if p_add_solve
- pressure_additional_solve!(
- pressure_solver,
- V,
- p,
- tₙ + Δt,
- setup,
- momentum_cache,
- F,
- f,
- Δp;
- bc_vectors,
- )
+ pressure!(psolver, V, p, tₙ + Δt, setup, momentum_cache, F, f, Δp; bc_vectors)
end
t = tₙ + Δtₙ
@@ -313,7 +300,7 @@ function step!(
create_stepper(
method;
setup,
- pressure_solver,
+ psolver,
bc_vectors,
V,
p,
diff --git a/src/time_steppers/step_explicit_runge_kutta.jl b/src/time_steppers/step_explicit_runge_kutta.jl
index bb25f37ac..6aa68fafc 100644
--- a/src/time_steppers/step_explicit_runge_kutta.jl
+++ b/src/time_steppers/step_explicit_runge_kutta.jl
@@ -1,198 +1,127 @@
-create_stepper(
- ::ExplicitRungeKuttaMethod;
- setup,
- pressure_solver,
- bc_vectors,
- V,
- p,
- t,
- n = 0,
-) = (; setup, pressure_solver, bc_vectors, V, p, t, n)
-
-function step(method::ExplicitRungeKuttaMethod, stepper, Δt)
- (; setup, pressure_solver, bc_vectors, V, p, t, n) = stepper
- (; grid, operators, boundary_conditions) = setup
- (; bc_unsteady) = boundary_conditions
- (; Ω) = grid
- (; G, M) = operators
- (; A, b, c, p_add_solve) = method
-
- T = typeof(Δt)
-
- # Update current solution (does not depend on previous step size)
- n += 1
- Vₙ = V
- pₙ = p
- tₙ = t
- Δtₙ = Δt
-
- # Number of stages
- nV = length(V)
- np = length(p)
+create_stepper(::ExplicitRungeKuttaMethod; setup, psolver, u, temp, t, n = 0) =
+ (; setup, psolver, u, temp, t, n)
+
+function timestep!(method::ExplicitRungeKuttaMethod, stepper, Δt; θ = nothing, cache)
+ (; setup, psolver, u, temp, t, n) = stepper
+ (; grid, closure_model, temperature) = setup
+ (; dimension, Iu) = grid
+ (; A, b, c) = method
+ (; u₀, ku, div, p, temp₀, ktemp, diff) = cache
+ D = dimension()
nstage = length(b)
+ m = closure_model
- # Reset RK arrays
- tᵢ = tₙ
- # kV = zeros(T, nV, 0)
- # kp = zeros(T, np, 0)
- kV = fill(V, 0)
-
- ## Start looping over stages
+ # Update current solution
+ t₀ = t
+ copyto!.(u₀, u)
+ isnothing(temp) || copyto!(temp₀, temp)
- # At i = 1 we calculate F₁, p₂ and u₂
- # ⋮
- # At i = s we calculate Fₛ, pₙ₊₁, and uₙ₊₁
for i = 1:nstage
- # Right-hand side for tᵢ based on current velocity field uₕ, vₕ at level i. This
- # includes force evaluation at tᵢ and pressure gradient. Boundary conditions will be
- # set through `get_bc_vectors` inside momentum. The pressure p is not important here,
- # it will be removed again in the next step
- F, ∇F = momentum(V, V, p, tᵢ, setup; bc_vectors)
-
- # Store right-hand side of stage i
- # Remove the -G*p contribution (but not y_p)
- kVᵢ = 1 ./ Ω .* (F + G * p)
- # kV = [kV kVᵢ]
- kV = [kV; [kVᵢ]]
-
- # Update velocity current stage by sum of Fᵢ's until this stage, weighted
- # with Butcher tableau coefficients. This gives uᵢ₊₁, and for i=s gives uᵢ₊₁
- # V = dot(kV * A[i, 1:i]
- V = sum(A[i, j] * kV[j] for j = 1:i)
-
- # Boundary conditions at tᵢ₊₁
- tᵢ = tₙ + c[i] * Δtₙ
- if bc_unsteady
- bc_vectors = get_bc_vectors(setup, tᵢ)
+ # Compute force at current stage i
+ apply_bc_u!(u, t, setup)
+ isnothing(temp) || apply_bc_temp!(temp, t, setup)
+ momentum!(ku[i], u, temp, t, setup)
+ if !isnothing(temp)
+ ktemp[i] .= 0
+ convection_diffusion_temp!(ktemp[i], u, temp, setup)
+ temperature.dodissipation && dissipation!(ktemp[i], diff, u, setup)
end
- (; yM) = bc_vectors
-
- # Divergence of intermediate velocity field
- f = (M * (Vₙ / Δtₙ + V) + yM / Δtₙ) / c[i]
- # Solve the Poisson equation, but not for the first step if the boundary conditions are steady
- if boundary_conditions.bc_unsteady || i > 1
- p = pressure_poisson(pressure_solver, f)
- else
- # Bc steady AND i = 1
- p = pₙ
- end
+ # Add closure term
+ isnothing(m) || map((k, m) -> k .+= m, ku[i], m(u, θ))
- Gp = G * p
+ # Intermediate time step
+ t = t₀ + c[i] * Δt
- # Update velocity current stage, which is now divergence free
- V = @. Vₙ + Δtₙ * (V - c[i] / Ω * Gp)
- end
+ # Apply stage forces
+ for α = 1:D
+ u[α] .= u₀[α]
+ for j = 1:i
+ @. u[α] += Δt * A[i, j] * ku[j][α]
+ # @. u[α][Iu[α]] += Δt * A[i, j] * ku[j][α][Iu[α]]
+ end
+ end
+ if !isnothing(temp)
+ temp .= temp₀
+ for j = 1:i
+ @. temp += Δt * A[i, j] * ktemp[j]
+ end
+ end
- # For steady bc we do an additional pressure solve
- # That saves a pressure solve for i = 1 in the next time step
- if !bc_unsteady || p_add_solve
- p = pressure_additional_solve(pressure_solver, V, p, tₙ + Δtₙ, setup; bc_vectors)
+ # Project stage u directly
+ # Make velocity divergence free at time t
+ apply_bc_u!(u, t, setup)
+ project!(u, setup; psolver, div, p)
end
- t = tₙ + Δtₙ
+ # This is redundant, but Neumann BC need to have _exact_ copies
+ # since we divide by an infinitely thin (eps(T)) volume width in the
+ # diffusion term
+ apply_bc_u!(u, t, setup)
+ isnothing(temp) || apply_bc_temp!(temp, t, setup)
- create_stepper(method; setup, pressure_solver, bc_vectors, V, p, t, n)
+ create_stepper(method; setup, psolver, u, temp, t, n = n + 1)
end
-function step!(method::ExplicitRungeKuttaMethod, stepper, Δt; cache, momentum_cache)
- (; setup, pressure_solver, bc_vectors, V, p, t, n) = stepper
- (; grid, operators, boundary_conditions) = setup
- (; bc_unsteady) = boundary_conditions
- (; Ω) = grid
- (; G, M) = operators
- (; A, b, c, p_add_solve) = method
- (; Vₙ, pₙ, kV, kp, Vtemp, Vtemp2, F, ∇F, Δp, f) = cache
+function timestep(method::ExplicitRungeKuttaMethod, stepper, Δt; θ = nothing)
+ (; setup, psolver, u, temp, t, n) = stepper
+ (; grid, closure_model, temperature) = setup
+ (; dimension) = grid
+ (; A, b, c) = method
+ D = dimension()
+ nstage = length(b)
+ m = closure_model
# Update current solution (does not depend on previous step size)
- n += 1
- Vₙ .= V
- pₙ .= p
- tₙ = t
- Δtₙ = Δt
+ t₀ = t
+ u₀ = u
+ ku = ()
+ ktemp = ()
- # Number of stages
- nstage = length(b)
-
- # # Reset RK arrays
- # kV .= 0
- # kp .= 0
+ for i = 1:nstage
+ # Compute force at current stage i
+ u = apply_bc_u(u, t, setup)
+ isnothing(temp) || (temp = apply_bc_temp(temp, t, setup))
+ F = momentum(u, temp, t, setup)
+ if !isnothing(temp)
+ Ftemp = convection_diffusion_temp(u, temp, setup)
+ temperature.dodissipation && (Ftemp += dissipation(u, setup))
+ end
- tᵢ = tₙ
+ # Add closure term
+ isnothing(m) || (F = F .+ m(u, θ))
- ## Start looping over stages
+ # Store right-hand side of stage i
+ ku = (ku..., F)
+ isnothing(temp) || (ktemp = (ktemp..., Ftemp))
- # At i = 1 we calculate F₁, p₂ and u₂
- # ⋮
- # At i = s we calculate Fₛ, pₙ₊₁, and uₙ₊₁
- for i = 1:nstage
- # Right-hand side for tᵢ based on current velocity field uₕ, vₕ at level i. This
- # includes force evaluation at tᵢ and pressure gradient. Boundary conditions will be
- # set through `get_bc_vectors` inside momentum. The pressure p is not important here,
- # it will be removed again in the next step
- momentum!(F, ∇F, V, V, p, tᵢ, setup, momentum_cache; bc_vectors)
+ # Intermediate time step
+ t = t₀ + c[i] * Δt
- # Store right-hand side of stage i
- # Remove the -G*p contribution (but not y_p)
- # kVᵢ = @view kV[:, i]
- kVᵢ = kV[i]
- mul!(kVᵢ, G, p)
- @. kVᵢ = 1 ./ Ω * (F + kVᵢ)
- # kVᵢ .= 1 ./ Ω .* (F + G * p)
-
- # Update velocity current stage by sum of Fᵢ's until this stage, weighted
- # with Butcher tableau coefficients. This gives uᵢ₊₁, and for i=s gives uᵢ₊₁
- # mul!(Vtemp, kV, A[i, :])
- Vtemp .= 0
+ # Apply stage forces
+ u = u₀
for j = 1:i
- Vtemp .= Vtemp .+ A[i, j] .* kV[j]
+ u = @. u + Δt * A[i, j] * ku[j]
+ # u = tupleadd(u, @.(Δt * A[i, j] * ku[j]))
end
-
- # Boundary conditions at tᵢ₊₁
- tᵢ = tₙ + c[i] * Δtₙ
- if bc_unsteady
- bc_vectors = get_bc_vectors(setup, tᵢ)
- end
- (; yM) = bc_vectors
-
- # Divergence of intermediate velocity field
- @. Vtemp2 = Vₙ / Δtₙ + Vtemp
- mul!(f, M, Vtemp2)
- @. f = (f + yM / Δtₙ) / c[i]
- # f = (M * (Vₙ / Δtₙ + Vtemp) + yM / Δtₙ) / c[i]
-
- # Solve the Poisson equation, but not for the first step if the boundary conditions are steady
- if boundary_conditions.bc_unsteady || i > 1
- pressure_poisson!(pressure_solver, p, f)
- else
- # Bc steady AND i = 1
- p .= pₙ
+ if !isnothing(temp)
+ temp = temp₀
+ for j = 1:i
+ temp = @. temp + Δt * A[i, j] * ktemp[j]
+ end
end
- mul!(Vtemp2, G, p)
-
- # Update velocity current stage, which is now divergence free
- @. V = Vₙ + Δtₙ * (Vtemp - c[i] / Ω * Vtemp2)
- end
-
- # For steady bc we do an additional pressure solve
- # That saves a pressure solve for i = 1 in the next time step
- if !bc_unsteady || p_add_solve
- pressure_additional_solve!(
- pressure_solver,
- V,
- p,
- tₙ + Δtₙ,
- setup,
- momentum_cache,
- F,
- f,
- Δp;
- bc_vectors,
- )
+ # Project stage u directly
+ # Make velocity divergence free at time t
+ u = apply_bc_u(u, t, setup)
+ u = project(u, setup; psolver)
end
- t = tₙ + Δtₙ
+ # This is redundant, but Neumann BC need to have _exact_ copies
+ # since we divide by an infinitely thin (eps(T)) volume width in the
+ # diffusion term
+ u = apply_bc_u(u, t, setup)
+ isnothing(temp) || (temp = apply_bc_temp(temp, t, setup))
- create_stepper(method; setup, pressure_solver, bc_vectors, V, p, t, n)
+ create_stepper(method; setup, psolver, u, temp, t, n = n + 1)
end
diff --git a/src/time_steppers/step_implicit_runge_kutta.jl b/src/time_steppers/step_implicit_runge_kutta.jl
index 64f17475b..ec08d148d 100644
--- a/src/time_steppers/step_implicit_runge_kutta.jl
+++ b/src/time_steppers/step_implicit_runge_kutta.jl
@@ -1,19 +1,11 @@
-create_stepper(
- ::ImplicitRungeKuttaMethod;
- setup,
- pressure_solver,
- bc_vectors,
- V,
- p,
- t,
- n = 0,
-) = (; setup, pressure_solver, bc_vectors, V, p, t, n)
+create_stepper(::ImplicitRungeKuttaMethod; setup, psolver, bc_vectors, V, p, t, n = 0) =
+ (; setup, psolver, bc_vectors, V, p, t, n)
-function step(method::ImplicitRungeKuttaMethod, stepper, Δt)
+function timestep(method::ImplicitRungeKuttaMethod, stepper, Δt)
# TODO: Implement out-of-place IRK
error()
- (; setup, pressure_solver, bc_vectors, V, p, t, n) = stepper
+ (; setup, psolver, bc_vectors, V, p, t, n) = stepper
(; grid, operators, boundary_conditions) = setup
(; bc_unsteady) = boundary_conditions
(; NV, Np, Ω) = grid
@@ -142,20 +134,13 @@ function step(method::ImplicitRungeKuttaMethod, stepper, Δt)
(; yM) = bc_vectors
f = 1 / Δtₙ .* (M * V .+ yM)
- p = pressure_poisson!(pressure_solver, f)
+ p = poisson!(psolver, f)
mul!(Gp, G, p)
V = @. V - Δtₙ / Ω * Gp
if p_add_solve
- p = pressure_additional_solve(
- pressure_solver,
- V,
- p,
- tₙ + Δtₙ,
- setup;
- bc_vectors,
- )
+ p = pressure(psolver, V, p, tₙ + Δtₙ, setup; bc_vectors)
else
# Standard method; take last pressure
p = pⱼ[(end-Np+1):end]
@@ -163,7 +148,7 @@ function step(method::ImplicitRungeKuttaMethod, stepper, Δt)
else
# For steady BC we do an additional pressure solve
# That saves a pressure solve for iter = 1 in the next time step
- # pressure_additional_solve!(pressure_solver, V, p, tₙ + Δtₙ, setup, momentum_cache, F, f, Δp)
+ # pressure!(psolver, V, p, tₙ + Δtₙ, setup, momentum_cache, F, f, Δp)
# Standard method; take pressure of last stage
p = pⱼ[(end-Np+1):end]
@@ -171,11 +156,11 @@ function step(method::ImplicitRungeKuttaMethod, stepper, Δt)
t = tₙ + Δtₙ
- create_stepper(method; setup, pressure_solver, bc_vectors, V, p, t, n)
+ create_stepper(method; setup, psolver, bc_vectors, V, p, t, n)
end
-function step!(method::ImplicitRungeKuttaMethod, stepper, Δt; cache, momentum_cache)
- (; setup, pressure_solver, bc_vectors, V, p, t, n) = stepper
+function timestep!(method::ImplicitRungeKuttaMethod, stepper, Δt; cache, momentum_cache)
+ (; setup, psolver, bc_vectors, V, p, t, n) = stepper
(; grid, operators, boundary_conditions) = setup
(; bc_unsteady) = boundary_conditions
(; NV, Np, Ω) = grid
@@ -372,24 +357,13 @@ function step!(method::ImplicitRungeKuttaMethod, stepper, Δt; cache, momentum_c
f .= yM
mul!(f, M, V, 1 / Δtₙ, 1 / Δtₙ)
# f .= 1 / Δtₙ .* (M * V .+ yM)
- pressure_poisson!(pressure_solver, p, f)
+ poisson!(psolver, p, f)
mul!(Gp, G, p)
@. V -= Δtₙ / Ω * Gp
if p_add_solve
- pressure_additional_solve!(
- pressure_solver,
- V,
- p,
- tₙ + Δtₙ,
- setup,
- momentum_cache,
- F,
- f,
- Δp;
- bc_vectors,
- )
+ pressure!(psolver, V, p, tₙ + Δtₙ, setup, momentum_cache, F, f, Δp; bc_vectors)
else
# Standard method; take last pressure
p .= pⱼ[(end-Np+1):end]
@@ -397,7 +371,7 @@ function step!(method::ImplicitRungeKuttaMethod, stepper, Δt; cache, momentum_c
else
# For steady BC we do an additional pressure solve
# That saves a pressure solve for iter = 1 in the next time step
- # pressure_additional_solve!(pressure_solver, V, p, tₙ + Δtₙ, setup, momentum_cache, F, f, Δp)
+ # pressure!(psolver, V, p, tₙ + Δtₙ, setup, momentum_cache, F, f, Δp)
# Standard method; take pressure of last stage
p .= pⱼ[(end-Np+1):end]
@@ -405,18 +379,14 @@ function step!(method::ImplicitRungeKuttaMethod, stepper, Δt; cache, momentum_c
t = tₙ + Δtₙ
- create_stepper(method; setup, pressure_solver, bc_vectors, V, p, t, n)
+ create_stepper(method; setup, psolver, bc_vectors, V, p, t, n)
end
-"""
- momentum_allstage(Vⱼ, ϕⱼ, pⱼ, tⱼ, setup; bc_vectors, nstage, get_jacobian = false)
-
-Call momentum for multiple `(Vⱼ, pⱼ)` pairs, as required in implicit RK methods.
-
-Non-mutating/allocating/out-of-place version.
-
-See also [`momentum_allstage!`](@ref).
-"""
+# Call momentum for multiple `(Vⱼ, pⱼ)` pairs, as required in implicit RK methods.
+#
+# Non-mutating/allocating/out-of-place version.
+#
+# See also [`momentum_allstage!`](@ref).
function momentum_allstage(Vⱼ, ϕⱼ, pⱼ, tⱼ, setup; bc_vectors, nstage, get_jacobian = false)
(; NV, Np) = setup.grid
T = eltype(Vⱼ)
@@ -446,28 +416,11 @@ function momentum_allstage(Vⱼ, ϕⱼ, pⱼ, tⱼ, setup; bc_vectors, nstage, g
Fⱼ, ∇Fⱼ
end
-"""
- momentum_allstage!(
- Fⱼ,
- ∇Fⱼ,
- Vⱼ,
- ϕⱼ,
- pⱼ,
- tⱼ,
- setup,
- cache,
- momentum_cache;
- bc_vectors,
- nstage,
- get_jacobian = false,
- )
-
-Call momentum for multiple `(V, p)` pairs, as required in implicit RK methods.
-
-Mutating/non-allocating/in-place version.
-
-See also [`momentum_allstage`](@ref).
-"""
+# Call momentum for multiple `(V, p)` pairs, as required in implicit RK methods.
+#
+# Mutating/non-allocating/in-place version.
+#
+# See also [`momentum_allstage`](@ref).
function momentum_allstage!(
Fⱼ,
∇Fⱼ,
diff --git a/src/time_steppers/step_one_leg.jl b/src/time_steppers/step_one_leg.jl
index a050340c6..3ed344897 100644
--- a/src/time_steppers/step_one_leg.jl
+++ b/src/time_steppers/step_one_leg.jl
@@ -1,213 +1,170 @@
create_stepper(
method::OneLegMethod;
setup,
- pressure_solver,
- bc_vectors,
- V,
- p,
+ psolver,
+ u,
t,
n = 0,
+ p = pressure(u, t, setup; psolver),
# For the first step, these are not used
- Vₙ = copy(V),
- pₙ = copy(p),
- tₙ = t,
-) = (; setup, pressure_solver, bc_vectors, V, p, t, n, Vₙ, pₙ, tₙ)
+ uold = copy.(u),
+ pold = copy(p),
+ told = t,
+) = (; setup, psolver, u, p, t, n, uold, pold, told)
-function step(method::OneLegMethod, stepper, Δt)
- (; setup, pressure_solver, bc_vectors, V, p, t, n, Vₙ, pₙ, tₙ) = stepper
+function timestep(method::OneLegMethod, stepper, Δt; θ = nothing)
+ (; setup, psolver, u, p, t, n, uold, pold, told) = stepper
(; p_add_solve, β, method_startup) = method
- (; grid, operators, boundary_conditions) = setup
- (; bc_unsteady) = boundary_conditions
- (; G, M) = operators
+ (; grid, boundary_conditions) = setup
(; Ω) = grid
+ T = typeof(Δt)
# One-leg requires state at previous time step, which is not available at
# the first iteration. Do one startup step instead
if n == 0
- stepper_startup =
- create_stepper(method_startup; setup, pressure_solver, bc_vectors, V, p, t)
- n += 1
- Vₙ = V
- pₙ = p
- tₙ = t
- (; V, p, t) = step(method_startup, stepper_startup, Δt)
- return create_stepper(
- method;
- setup,
- pressure_solver,
- bc_vectors,
- V,
- p,
- t,
- n,
- Vₙ,
- pₙ,
- tₙ,
- )
+ stepper_startup = create_stepper(method_startup; setup, psolver, u, t)
+ (; u, t, n) = timestep(method_startup, stepper_startup, Δt)
+ p = pressure(u, t, setup; psolver)
+ return create_stepper(method; setup, psolver, u, p, t, n, uold, pold, told)
end
- # Update current solution
- Δtₙ₋₁ = t - tₙ
- n += 1
- Vₙ₋₁ = Vₙ
- pₙ₋₁ = pₙ
- Vₙ = V
- pₙ = p
- tₙ = t
- Δtₙ = Δt
-
# One-leg requires fixed time step
- @assert Δtₙ ≈ Δtₙ₋₁
+ @assert Δt ≈ t - told
# Intermediate ("offstep") velocities
- t = tₙ + β * Δtₙ
- V = @. (1 + β) * Vₙ - β * Vₙ₋₁
- p = @. (1 + β) * pₙ - β * pₙ₋₁
+ tβ = t + β * Δt
+ uβ = ntuple(α -> @.((1 + β) * u[α] - β * uold[α]), length(u))
+ pβ = @. (1 + β) * p - β * pold
# Right-hand side of the momentum equation
- F, = momentum(V, V, p, t, setup; bc_vectors)
+ F = momentum(uβ, tβ, setup)
+ G = pressuregradient(pβ, setup)
# Take a time step with this right-hand side, this gives an intermediate velocity field
# (not divergence free)
- V = @. (2β * Vₙ - (β - 1 // 2) * Vₙ₋₁ + Δtₙ / Ω * F) / (β + 1 // 2)
+ unew = ntuple(
+ α -> @.(
+ (2β * u[α] - (β - T(1) / 2) * uold[α] + Δt * (F[α] - G[α])) / (β + T(1) / 2)
+ ),
+ length(u),
+ )
# To make the velocity field uₙ₊₁ at tₙ₊₁ divergence-free we need the boundary
# conditions at tₙ₊₁
- if bc_unsteady
- bc_vectors = get_bc_vectors(setup, tₙ + Δtₙ)
- end
- (; yM) = bc_vectors
+ unew = apply_bc_u(unew, t + Δt, setup)
# Adapt time step for pressure calculation
- Δtᵦ = Δtₙ / (β + 1 // 2)
+ Δtᵦ = Δt / (β + T(1) / 2)
# Divergence of intermediate velocity field
- f = (M * V + yM) / Δtᵦ
+ div = divergence(unew, setup) / Δtᵦ
+ div = @. div * Ω
# Solve the Poisson equation for the pressure
- Δp = pressure_poisson(pressure_solver, f)
- GΔp = G * Δp
+ Δp = poisson(psolver, div)
+ Δp = apply_bc_p(Δp, t + Δtᵦ, setup)
+ GΔp = pressuregradient(Δp, setup)
# Update velocity field
- V = @. V - Δtᵦ / Ω * GΔp
+ unew = ntuple(α -> @.(unew[α] - Δtᵦ * GΔp[α]), length(u))
+ unew = apply_bc_u(unew, t + Δt, setup)
# Update pressure (second order)
- p = @. 2pₙ - pₙ₋₁ + 4 // 3 * Δp
+ pnew = @. 2 * p - pold + T(4) / 3 * Δp
+ pnew = apply_bc_p(pnew, t + Δt, setup)
# Alternatively, do an additional Poisson solve
if p_add_solve
- p = pressure_additional_solve(pressure_solver, V, p, tₙ + Δtₙ, setup; bc_vectors)
+ pnew = pressure(unew, t + Δt, setup; psolver)
end
- t = tₙ + Δtₙ
+ told = t
+ pold = p
+ uold = u
+ t = t + Δt
+ p = pnew
+ u = unew
- create_stepper(method; setup, pressure_solver, bc_vectors, V, p, t, n, Vₙ, pₙ, tₙ)
+ create_stepper(method; setup, psolver, u, p, t, n, uold, pold, told)
end
-function step!(method::OneLegMethod, stepper, Δt; cache, momentum_cache)
- (; setup, pressure_solver, bc_vectors, n, V, p, t, Vₙ, pₙ, tₙ) = stepper
+function timestep!(method::OneLegMethod, stepper, Δt; θ = nothing, cache)
+ (; setup, psolver, u, p, t, n, uold, pold, told) = stepper
(; p_add_solve, β, method_startup) = method
- (; grid, operators, boundary_conditions) = setup
- (; bc_unsteady) = boundary_conditions
- (; G, M) = operators
+ (; grid, boundary_conditions) = setup
(; Ω) = grid
- (; Vₙ₋₁, pₙ₋₁, F, f, Δp, GΔp) = cache
+ (; unew, pnew, div, F, Δp) = cache
+ T = typeof(Δt)
# One-leg requires state at previous time step, which is not available at
# the first iteration. Do one startup step instead
if n == 0
- stepper_startup =
- create_stepper(method_startup; setup, pressure_solver, bc_vectors, V, p, t)
- n += 1
- Vₙ = V
- pₙ = p
- tₙ = t
-
- # Note: We do one out-of-place step here, with a few allocations
- (; V, p, t) = step(method_startup, stepper_startup, Δt)
- return create_stepper(
- method;
- setup,
- pressure_solver,
- bc_vectors,
- V,
- p,
- t,
- n,
- Vₙ,
- pₙ,
- tₙ,
- )
+ stepper_startup = create_stepper(method_startup; setup, psolver, u, t)
+ (; u, t, n) = timestep(method_startup, stepper_startup, Δt)
+ pressure!(p, u, temp, t, setup; psolver, F, div)
+ return create_stepper(method; setup, psolver, u, p, t, n, uold, pold, told)
end
- # Update current solution
- Δtₙ₋₁ = t - tₙ
- n += 1
- Vₙ₋₁ .= Vₙ
- pₙ₋₁ .= pₙ
- Vₙ .= V
- pₙ .= p
- tₙ = t
- Δtₙ = Δt
-
# One-leg requires fixed time step
- @assert Δtₙ ≈ Δtₙ₋₁
+ @assert Δt ≈ t - told
# Intermediate ("offstep") velocities
- t = tₙ + β * Δtₙ
- @. V = (1 + β) * Vₙ - β * Vₙ₋₁
- @. p = (1 + β) * pₙ - β * pₙ₋₁
+ tnew = t + β * Δt
+ for α = 1:length(u)
+ @. unew[α] = (1 + β) * u[α] - β * uold[α]
+ end
+ @. pnew = (1 + β) * p - β * pold
# Right-hand side of the momentum equation
- momentum!(F, nothing, V, V, p, t, setup, momentum_cache)
+ momentum!(F, unew, tnew, setup)
+ applypressure!(F, pnew, setup)
# Take a time step with this right-hand side, this gives an intermediate velocity field
# (not divergence free)
- @. V = (2β * Vₙ - (β - 1 // 2) * Vₙ₋₁ + Δtₙ / Ω * F) / (β + 1 // 2)
+ for α = 1:length(u)
+ @. unew[α] = 2β * u[α] - (β - T(1) / 2) * uold[α] + Δt * F[α] / (β + T(1) / 2)
+ end
# To make the velocity field uₙ₊₁ at tₙ₊₁ divergence-free we need the boundary
# conditions at tₙ₊₁
- if bc_unsteady
- bc_vectors = get_bc_vectors(setup, tₙ + Δtₙ)
- end
- (; yM) = bc_vectors
+ apply_bc_u!(unew, t + Δt, setup)
# Adapt time step for pressure calculation
- Δtᵦ = Δtₙ / (β + 1 // 2)
+ Δtᵦ = Δt / (β + T(1) / 2)
# Divergence of intermediate velocity field
- f .= yM
- mul!(f, M, V, 1 / Δtᵦ, 1 / Δtᵦ)
- # f .= (M * V + yM) / Δtᵦ
+ divergence!(div, unew, setup)
+ @. div *= Ω
# Solve the Poisson equation for the pressure
- pressure_poisson!(pressure_solver, Δp, f)
- mul!(GΔp, G, Δp)
+ poisson!(psolver, Δp, div)
+ apply_bc_p!(Δp, t + Δtᵦ, setup)
# Update velocity field
- @. V -= Δtᵦ / Ω * GΔp
+ applypressure!(unew, Δp, setup)
+ apply_bc_u!(unew, t + Δt, setup)
# Update pressure (second order)
- @. p = 2pₙ - pₙ₋₁ + 4 // 3 * Δp
+ @. pnew = 2 * p - pold + T(4) / 3 * Δp / Δtᵦ
+ apply_bc_p!(pnew, t + Δt, setup)
# Alternatively, do an additional Poisson solve
if p_add_solve
- pressure_additional_solve!(
- pressure_solver,
- V,
- p,
- tₙ + Δtₙ,
- setup,
- momentum_cache,
- F,
- f,
- Δp;
- bc_vectors,
- )
+ pressure!(pnew, unew, tempnew, t + Δt, setup; psolver, F, div)
end
- t = tₙ + Δtₙ
+ n += 1
+ told = t
+ pold .= p
+ for α = 1:length(u)
+ uold[α] .= u[α]
+ end
+ t = t + Δt
+ p .= pnew
+ for α = 1:length(u)
+ u[α] .= unew[α]
+ end
- create_stepper(method; setup, pressure_solver, bc_vectors, V, p, t, n, Vₙ, pₙ, tₙ)
+ create_stepper(method; setup, psolver, u, p, t, n, uold, pold, told)
end
diff --git a/src/time_steppers/tableaux.jl b/src/time_steppers/tableaux.jl
index e9df37667..ba9bbc49e 100644
--- a/src/time_steppers/tableaux.jl
+++ b/src/time_steppers/tableaux.jl
@@ -1,9 +1,38 @@
"""
+ RKMethods
+
Set up Butcher arrays `A`, `b`, and `c`, as well as and SSP coefficient `r`.
For families of methods, optional input `s` is the number of stages.
Original (MATLAB) by David Ketcheson, extended by Benjamin Sanderse.
"""
+module RKMethods
+
+using IncompressibleNavierStokes: runge_kutta_method
+
+# Explicit Methods
+export FE11, SSP22, SSP42, SSP33, SSP43, SSP104, rSSPs2, rSSPs3, Wray3, RK56, DOPRI6
+
+# Implicit Methods
+export BE11, SDIRK34, ISSPm2, ISSPs3
+
+# Half explicit methods
+export HEM3, HEM3BS, HEM5
+
+# Classical Methods
+export GL1, GL2, GL3, RIA1, RIA2, RIA3, RIIA1, RIIA2, RIIA3, LIIIA2, LIIIA3
+
+# Chebyshev methods
+export CHDIRK3, CHCONS3, CHC3, CHC5
+
+# Miscellaneous Methods
+export Mid22, MTE22, CN22, Heun33, RK33C2, RK33P2, RK44, RK44C2, RK44C23, RK44P2
+
+# DSRK Methods
+export DSso2, DSRK2, DSRK3
+
+# "Non-SSP" Methods of Wong & Spiteri
+export NSSP21, NSSP32, NSSP33, NSSP53
# Default SSP coefficient
# r = 0
@@ -807,3 +836,5 @@ function NSSP53(; kwargs...)
c = [0, 1 // 7, 3 // 16, 1 // 3, 2 // 3]
runge_kutta_method(A, b, c, r; kwargs...)
end
+
+end
diff --git a/src/time_steppers/time_stepper_caches.jl b/src/time_steppers/time_stepper_caches.jl
index b8b4716af..6c17cfa8d 100644
--- a/src/time_steppers/time_stepper_caches.jl
+++ b/src/time_steppers/time_stepper_caches.jl
@@ -7,55 +7,47 @@ Get time stepper cache for the given ODE method.
function ode_method_cache end
function ode_method_cache(::AdamsBashforthCrankNicolsonMethod, setup, V, p)
- cₙ = similar(V)
- cₙ₋₁ = similar(V)
- F = similar(V)
- f = similar(p)
- Δp = similar(p)
- Rr = similar(V)
- b = similar(V)
- bₙ = similar(V)
- bₙ₊₁ = similar(V)
- yDiffₙ = similar(V)
- yDiffₙ₊₁ = similar(V)
- Gpₙ = similar(V)
-
- (; cₙ, cₙ₋₁, F, f, Δp, Rr, b, bₙ, bₙ₊₁, yDiffₙ, yDiffₙ₊₁, Gpₙ)
+ c₀ = zero(V)
+ c₋₁ = zero(V)
+ F = zero(V)
+ f = zero(p)
+ Δp = zero(p)
+ Rr = zero(V)
+ b = zero(V)
+ b₀ = zero(V)
+ b₁ = zero(V)
+ yDiff₀ = zero(V)
+ yDiff₁ = zero(V)
+ Gp₀ = zero(V)
+
+ (; c₀, c₋₁, F, f, Δp, Rr, b, b₀, b₁, yDiff₀, yDiff₁, Gp₀)
end
-function ode_method_cache(::OneLegMethod{T}, setup, V, p) where {T}
- (; NV, Np) = setup.grid
- Vₙ₋₁ = similar(V)
- pₙ₋₁ = similar(p)
- F = similar(V)
- f = similar(p)
- Δp = similar(p)
- GΔp = similar(V)
- (; Vₙ₋₁, pₙ₋₁, F, f, Δp, GΔp)
+function ode_method_cache(::OneLegMethod{T}, setup, u) where {T}
+ unew = zero.(u)
+ pnew = zero(u[1])
+ div = zero(u[1])
+ F = zero.(u)
+ Δp = zero(u[1])
+ (; unew, pnew, F, div, Δp)
end
-function ode_method_cache(method::ExplicitRungeKuttaMethod{T}, setup, V, p) where {T}
- (; NV, Np) = setup.grid
-
- Vₙ = similar(V)
- pₙ = similar(p)
-
+function ode_method_cache(method::ExplicitRungeKuttaMethod, setup, u, temp)
+ u₀ = zero.(u)
ns = nstage(method)
-
- # kV = zeros(T, NV, ns)
- # kp = zeros(T, Np, ns)
-
- kV = [similar(V) for i = 1:ns]
- kp = [similar(p) for i = 1:ns]
-
- Vtemp = similar(V)
- Vtemp2 = similar(V)
- F = similar(V)
- ∇F = spzeros(T, NV, NV)
- f = similar(p)
- Δp = similar(p)
-
- (; Vₙ, pₙ, kV, kp, Vtemp, Vtemp2, F, ∇F, f, Δp)
+ ku = [zero.(u) for i = 1:ns]
+ div = zero(u[1])
+ p = zero(u[1])
+ if isnothing(temp)
+ temp₀ = nothing
+ ktemp = nothing
+ diff = nothing
+ else
+ temp₀ = copy(temp)
+ ktemp = [copy(temp) for i = 1:ns]
+ diff = zero.(u)
+ end
+ (; u₀, ku, div, p, temp₀, ktemp, diff)
end
function ode_method_cache(method::ImplicitRungeKuttaMethod{T}, setup, V, p) where {T}
@@ -63,8 +55,8 @@ function ode_method_cache(method::ImplicitRungeKuttaMethod{T}, setup, V, p) wher
(; G, M) = setup.operators
(; A, b, c) = method
- Vₙ = similar(V)
- pₙ = similar(p)
+ Vₙ = zero(V)
+ pₙ = zero(p)
# Number of stages
s = length(b)
@@ -85,11 +77,11 @@ function ode_method_cache(method::ImplicitRungeKuttaMethod{T}, setup, V, p) wher
fⱼ = zeros(T, s * (NV + Np))
- F = similar(V)
+ F = zero(V)
∇F = spzeros(T, NV, NV)
- f = similar(p)
- Δp = similar(p)
- Gp = similar(V)
+ f = zero(p)
+ Δp = zero(p)
+ Gp = zero(V)
# Gradient operator (could also use 1 instead of c and later scale the pressure)
Gtot = kron(A, G)
diff --git a/src/utils/filter_convection.jl b/src/utils/filter_convection.jl
deleted file mode 100644
index 632fbc58d..000000000
--- a/src/utils/filter_convection.jl
+++ /dev/null
@@ -1,28 +0,0 @@
-"""
- filter_convection(u, diff_matrix, bc, α)
-
-Construct filter for convective terms.
-
-Non-mutating/allocating/out-of-place version.
-
-See also [`filter_convection!`](@ref).
-"""
-function filter_convection(u, diff_matrix, bc, α)
- u + α * (diff_matrix * u + bc)
-end
-
-"""
- filter_convection!(ū, u, diff_matrix, bc, α)
-
-Construct filter for convective terms.
-
-Mutating/non-allocating/in-place version.
-
-See also [`filter_convection`](@ref).
-"""
-function filter_convection!(ū, u, diff_matrix, bc, α)
- # ū = u + α * (diff_matrix * u + bc)
- mul!(ū, diff_matrix, u)
- @. ū = u + α * (ū + bc)
- ū
-end
diff --git a/src/utils/get_lims.jl b/src/utils/get_lims.jl
index cc7f1a730..8e67c4470 100644
--- a/src/utils/get_lims.jl
+++ b/src/utils/get_lims.jl
@@ -6,8 +6,9 @@ deviation (``\\mu \\pm n \\sigma``). If `x` is constant, a margin of `1e-4` is e
plotting functions that require a certain range.
"""
function get_lims(x, n = 1.5)
+ T = eltype(x)
μ = mean(x)
σ = std(x)
- ≈(μ + σ, μ; rtol = 1e-8, atol = 1e-8) && (σ = 1e-4)
+ ≈(μ + σ, μ; rtol = sqrt(eps(T)), atol = sqrt(eps(T))) && (σ = sqrt(sqrt(eps(T))))
(μ - n * σ, μ + n * σ)
end
diff --git a/src/postprocess/plot_grid.jl b/src/utils/plotgrid.jl
similarity index 90%
rename from src/postprocess/plot_grid.jl
rename to src/utils/plotgrid.jl
index 8d6717db9..cc03cdf4c 100644
--- a/src/postprocess/plot_grid.jl
+++ b/src/utils/plotgrid.jl
@@ -1,20 +1,18 @@
"""
- plot_grid(x, y)
- plot_grid(x, y, z)
- plot_grid(grid)
+ plotgrid(x...)
Plot nonuniform Cartesian grid.
"""
-function plot_grid end
+function plotgrid end
-plot_grid(x, y) = wireframe(
+plotgrid(x, y) = wireframe(
x,
y,
zeros(eltype(x), length(x), length(y));
axis = (; aspect = DataAspect(), xlabel = "x", ylabel = "y"),
)
-function plot_grid(x, y, z)
+function plotgrid(x, y, z)
nx, ny, nz = length(x), length(y), length(z)
T = eltype(x)
diff --git a/src/utils/save_vtk.jl b/src/utils/save_vtk.jl
new file mode 100644
index 000000000..0456e6cc2
--- /dev/null
+++ b/src/utils/save_vtk.jl
@@ -0,0 +1,39 @@
+"""
+ save_vtk(setup, u, filename = "output/solution"; fieldnames = [:velocity], psolver)
+
+Save velocity and pressure field to a VTK file.
+
+In the case of a 2D setup, the velocity field is saved as a 3D vector with a
+z-component of zero, as this seems to be preferred by ParaView.
+"""
+function save_vtk(
+ setup,
+ u,
+ t,
+ filename = "output/solution";
+ fieldnames = [:velocity],
+ psolver = default_psolver(setup),
+)
+ parts = split(filename, "/")
+ path = join(parts[1:end-1], "/")
+ isdir(path) || mkpath(path)
+ (; grid) = setup
+ (; dimension, xp) = grid
+ D = dimension()
+ xp = Array.(xp)
+ vtk_grid(filename, xp...) do vtk
+ up = interpolate_u_p(u, setup)
+ ωp = interpolate_ω_p(vorticity(u, setup), setup)
+ if D == 2
+ # ParaView prefers 3D vectors. Add zero z-component.
+ up3 = zero(up[1])
+ up = (up..., up3)
+ ωp = Array(ωp)
+ else
+ ωp = Array.(ωp)
+ end
+ vtk["velocity"] = Array.(up)
+ :pressure in fieldnames && (vtk["pressure"] = Array(pressure(u, t, setup; psolver)))
+ vtk["vorticity"] = ωp
+ end
+end
diff --git a/src/utils/spectral_stuff.jl b/src/utils/spectral_stuff.jl
new file mode 100644
index 000000000..471c32d50
--- /dev/null
+++ b/src/utils/spectral_stuff.jl
@@ -0,0 +1,88 @@
+function spectral_stuff(setup; npoint = 100, a = typeof(setup.Re)(1 + sqrt(5)) / 2)
+ (; dimension, xp, Ip) = setup.grid
+ T = eltype(xp[1])
+ D = dimension()
+
+ K = size(Ip) .÷ 2
+ k = zeros(T, K)
+ for α = 1:D
+ kα =
+ reshape(0:K[α]-1, ntuple(Returns(1), α - 1)..., :, ntuple(Returns(1), D - α)...)
+ k .+= kα .^ 2
+ end
+ k .= sqrt.(k)
+ k = reshape(k, :)
+
+ # Sum or average wavenumbers between k and k+1
+ kmax = minimum(K) - 1
+ isort = sortperm(k)
+ ksort = k[isort]
+ ia = zeros(Int, 0)
+ ib = zeros(Int, 0)
+ vals = zeros(T, 0)
+
+ # Output query points (evenly log-spaced, but only integer wavenumbers)
+ # logκ = LinRange(T(0), log(T(kmax) - 1), npoint)
+ logκ = LinRange(T(0), log(T(kmax) / a), npoint)
+ # logκ = LinRange(log(a), log(T(kmax) / a), npoint)
+ # logκ = LinRange(T(0), log(T(kmax)), npoint)
+ κ = exp.(logκ)
+ κ = sort(unique(round.(Int, κ)))
+ npoint = length(κ)
+
+ for i = 1:npoint
+ jstart = findfirst(≥(κ[i] / a), ksort)
+ jstop = findfirst(≥(κ[i] * a), ksort)
+ # jstart = findfirst(≥(κ[i] - T(1.01)), ksort)
+ # jstop = findfirst(≥(κ[i] + T(1.01)), ksort)
+ isnothing(jstop) && (jstop = length(ksort) + 1)
+ jstop -= 1
+ nk = jstop - jstart + 1
+ append!(ia, fill(i, nk))
+ append!(ib, isort[jstart:jstop])
+ append!(vals, fill(T(1), nk))
+ end
+ IntArray = typeof(similar(xp[1], Int, 0))
+ TArray = typeof(similar(xp[1], 0))
+ ia = adapt(IntArray, ia)
+ ib = adapt(IntArray, ib)
+ vals = adapt(TArray, vals)
+ A = sparse(ia, ib, vals, npoint, length(k))
+
+ (; A, κ, K)
+end
+
+function get_spectrum(setup; npoint = 100, a = typeof(e.setup.Re)(1 + sqrt(5)) / 2)
+ (; dimension, xp, Ip) = setup.grid
+ T = eltype(xp[1])
+ D = dimension()
+
+ @assert all(==(size(Ip, 1)), size(Ip))
+
+ K = size(Ip, 1) .÷ 2
+ kmax = K - 1
+ k = ntuple(
+ i -> reshape(0:kmax, ntuple(Returns(1), i - 1)..., :, ntuple(Returns(1), D - i)...),
+ D,
+ )
+
+ # Output query points (evenly log-spaced, but only integer wavenumbers)
+ logκ = LinRange(T(0), log(T(kmax) / a), npoint)
+ κ = exp.(logκ)
+ κ = sort(unique(round.(Int, κ)))
+ npoint = length(κ)
+
+ masks = map(κ) do κ
+ if D == 2
+ @. (κ / a)^2 ≤ k[1]^2 + k[2]^2 < (κ * a)^2
+ elseif D == 3
+ @. (κ / a)^2 ≤ k[1]^2 + k[2]^2 + k[3]^2 < (κ * a)^2
+ else
+ error("Not implemented")
+ end
+ end
+
+ BoolArray = typeof(similar(xp[1], Bool, ntuple(Returns(0), D)...))
+ masks = adapt.(BoolArray, masks)
+ (; κ, masks, K)
+end
diff --git a/test/Project.toml b/test/Project.toml
new file mode 100644
index 000000000..47bb34b07
--- /dev/null
+++ b/test/Project.toml
@@ -0,0 +1,10 @@
+[deps]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
diff --git a/test/chainrules.jl b/test/chainrules.jl
new file mode 100644
index 000000000..eff75485b
--- /dev/null
+++ b/test/chainrules.jl
@@ -0,0 +1,72 @@
+# @testset "Chain rules boundary conditions" begin
+# T = Float64
+# Re = T(1_000)
+# n = 8
+# boundary_conditions = ((DirichletBC(), PressureBC()), (PeriodicBC(), PeriodicBC()))
+# lims = T(0), T(1)
+# x = stretched_grid(lims..., n, 1.2), cosine_grid(lims..., n)
+# setup = Setup(x...; Re,
+# boundary_conditions,
+# )
+# u = random_field(setup, T(0))
+# randn!.(u)
+# p = randn!(similar(u[1]))
+# test_rrule(apply_bc_u, u, T(0) ⊢ NoTangent(), setup ⊢ NoTangent())
+# test_rrule(apply_bc_p, p, T(0) ⊢ NoTangent(), setup ⊢ NoTangent())
+# end;
+
+# Test chain rule correctness by comparing with finite differences
+testchainrules(dim) = @testset "Chain rules $(dim())D" begin
+ @info "Testing chain rules in $(dim())D"
+
+ # Setup
+ D = dim()
+ T = Float64
+ Re = T(1_000)
+ n = if D == 2
+ 8
+ elseif D == 3
+ # 4^3 = 64 grid points
+ # 3*64 = 192 velocity components
+ # 192^2 = 36864 fininite difference pairs in convection/diffusion
+ # TODO: Check if `test_rrule` computes all combinations or only a subset
+ 4
+ end
+ lims = T(0), T(1)
+ x = if D == 2
+ stretched_grid(lims..., n, 1.2), cosine_grid(lims..., n)
+ elseif D == 3
+ stretched_grid(lims..., n, 1.2), cosine_grid(lims..., n), cosine_grid(lims..., n)
+ end
+ setup = Setup(x...; Re)
+ psolver = default_psolver(setup)
+ u = random_field(setup, T(0); psolver)
+ randn!.(u)
+ p = randn!(similar(u[1]))
+ @testset "Boundary conditions" begin
+ test_rrule(apply_bc_u, u, T(0) ⊢ NoTangent(), setup ⊢ NoTangent())
+ test_rrule(apply_bc_p, p, T(0) ⊢ NoTangent(), setup ⊢ NoTangent())
+ end
+ @testset "Divergence" begin
+ test_rrule(divergence, u, setup ⊢ NoTangent())
+ end
+ @testset "Pressure gradient" begin
+ test_rrule(pressuregradient, p, setup ⊢ NoTangent())
+ end
+ @testset "Poisson" begin
+ test_rrule(poisson, psolver ⊢ NoTangent(), p)
+ end
+ @testset "Convection" begin
+ test_rrule(convection, u, setup ⊢ NoTangent())
+ end
+ @testset "Diffusion" begin
+ test_rrule(diffusion, u, setup ⊢ NoTangent())
+ end
+ @testset "Bodyforce" begin
+ @test_broken 1 == 2 # Just to identify location for broken rrule test
+ # test_rrule(bodyforce, u, T(0) ⊢ NoTangent(), setup ⊢ NoTangent())
+ end
+end
+
+testchainrules(IncompressibleNavierStokes.Dimension(2))
+testchainrules(IncompressibleNavierStokes.Dimension(3));
diff --git a/test/models.jl b/test/models.jl
index 5f23e7473..15f911bc8 100644
--- a/test/models.jl
+++ b/test/models.jl
@@ -23,10 +23,10 @@
# Viscosity models
T = Float64
Re = 1000.0
- lam = LaminarModel(; Re)
- ml = MixingLengthModel(; Re, lm = Δ)
- smag = SmagorinskyModel(; Re)
- qr = QRModel(; Re)
+ lam = LaminarModel()
+ ml = MixingLengthModel(; lm = Δ)
+ smag = SmagorinskyModel()
+ qr = QRModel()
# Convection models
noreg = NoRegConvectionModel()
@@ -46,9 +46,9 @@
for (viscosity_model, convection_model) in models
@testset "$(typeof(viscosity_model)) $(typeof(convection_model))" begin
@info "Testing $(typeof(viscosity_model)) and $(typeof(convection_model))"
- setup = Setup(x, y; viscosity_model, convection_model, u_bc, v_bc, bc_type)
+ setup = Setup(x, y; Re, viscosity_model, convection_model, u_bc, v_bc, bc_type)
- V₀, p₀ = create_initial_conditions(
+ V = create_initial_conditions(
setup,
initial_velocity_u,
initial_velocity_v,
@@ -60,7 +60,7 @@
broken = convection_model isa Union{C2ConvectionModel,C4ConvectionModel}
@test sum(abs, V) / length(V) < lid_vel broken = broken
- V, p, outputs = solve_unsteady(setup, V₀, p₀, tlims; Δt = 0.01)
+ (; u, t), outputs = solve_unsteady(setup, V₀, tlims; Δt = 0.01)
# Check that the average velocity is smaller than the lid velocity
broken = convection_model isa Union{C2ConvectionModel,C4ConvectionModel}
@@ -73,7 +73,7 @@
for (viscosity_model, convection_model) in models
@testset "$(typeof(viscosity_model)) $(typeof(convection_model))" begin
@info "Testing $(typeof(viscosity_model)) and $(typeof(convection_model))"
- setup = Setup(x, y; viscosity_model, convection_model, u_bc, v_bc, bc_type)
+ setup = Setup(x, y; Re, viscosity_model, convection_model, u_bc, v_bc, bc_type)
end
end
end
diff --git a/test/operators.jl b/test/operators.jl
new file mode 100644
index 000000000..1bca2e7c4
--- /dev/null
+++ b/test/operators.jl
@@ -0,0 +1,146 @@
+testops(dim) = @testset "Operators $(dim())D" begin
+ @info "Testing operators in $(dim())D"
+
+ # Setup
+ D = dim()
+ T = Float64
+ Re = T(1_000)
+ n = 16
+ lims = T(0), T(1)
+ x = if D == 2
+ stretched_grid(lims..., n, 1.2), cosine_grid(lims..., n)
+ elseif D == 3
+ stretched_grid(lims..., n, 1.2), cosine_grid(lims..., n), cosine_grid(lims..., n)
+ end
+ setup = Setup(x...; Re)
+ u = random_field(setup, T(0))
+ (; Iu, Ip, Ω) = setup.grid
+
+ @testset "Divergence" begin
+ div = divergence(u, setup)
+ @test div isa Array{T}
+ @test all(!isnan, div)
+ end
+
+ @testset "Pressure gradient" begin
+ v = randn!.(similar.(u))
+ p = randn!(similar(u[1]))
+ v = apply_bc_u(v, T(0), setup)
+ p = apply_bc_p(p, T(0), setup)
+ Dv = divergence(v, setup)
+ Gp = pressuregradient(p, setup)
+ pDv = sum((p.*Ω.*Dv)[Ip])
+ vGp = if D == 2
+ vGpx = v[1] .* setup.grid.Δu[1] .* setup.grid.Δ[2]' .* Gp[1]
+ vGpy = v[2] .* setup.grid.Δ[1] .* setup.grid.Δu[2]' .* Gp[2]
+ sum(vGpx[Iu[1]]) + sum(vGpy[Iu[2]])
+ elseif D == 3
+ vGpx =
+ v[1] .* setup.grid.Δu[1] .* reshape(setup.grid.Δ[2], 1, :) .*
+ reshape(setup.grid.Δ[3], 1, 1, :) .* Gp[1]
+ vGpy =
+ v[2] .* setup.grid.Δ[1] .* reshape(setup.grid.Δu[2], 1, :) .*
+ reshape(setup.grid.Δ[3], 1, 1, :) .* Gp[2]
+ vGpz =
+ v[3] .* setup.grid.Δ[1] .* reshape(setup.grid.Δ[2], 1, :) .*
+ reshape(setup.grid.Δu[3], 1, 1, :) .* Gp[3]
+ sum(vGpx[Iu[1]]) + sum(vGpy[Iu[2]]) + sum(vGpz[Iu[3]])
+ end
+ @test Gp isa Tuple
+ @test Gp[1] isa Array{T}
+ @test pDv ≈ -vGp # Check that D = -G'
+ end
+
+ @testset "Laplacian" begin
+ p = randn!(similar(u[1]))
+ p = apply_bc_p(p, T(0), setup)
+ Lp = laplacian(p, setup)
+ @test Lp isa Array{T}
+ @test sum((p.*Ω.*Lp)[Ip]) ≤ 0 # Check negativity
+ L = laplacian_mat(setup)
+ @test L isa SparseMatrixCSC
+ @test norm((Lp)[Ip][:] - (L * p[Ip][:])) ≈ 0 atol = 1e-12
+ end
+
+ @testset "Convection" begin
+ c = convection(u, setup)
+ uCu = if D == 2
+ uCux = u[1] .* setup.grid.Δu[1] .* setup.grid.Δ[2]' .* c[1]
+ uCuy = u[2] .* setup.grid.Δ[1] .* setup.grid.Δu[2]' .* c[2]
+ sum(uCux[Iu[1]]) + sum(uCuy[Iu[2]])
+ elseif D == 3
+ uCux =
+ u[1] .* setup.grid.Δu[1] .* reshape(setup.grid.Δ[2], 1, :) .*
+ reshape(setup.grid.Δ[3], 1, 1, :) .* c[1]
+ uCuy =
+ u[2] .* setup.grid.Δ[1] .* reshape(setup.grid.Δu[2], 1, :) .*
+ reshape(setup.grid.Δ[3], 1, 1, :) .* c[2]
+ uCuz =
+ u[3] .* setup.grid.Δ[1] .* reshape(setup.grid.Δ[2], 1, :) .*
+ reshape(setup.grid.Δu[3], 1, 1, :) .* c[3]
+ sum(uCux[Iu[1]]) + sum(uCuy[Iu[2]]) + sum(uCuz[Iu[3]])
+ end
+ @test c isa Tuple
+ @test c[1] isa Array{T}
+ @test uCu ≈ 0 atol = 1e-12 # Check skew-symmetry
+ end
+
+ @testset "Diffusion" begin
+ d = diffusion(u, setup)
+ uDu = if D == 2
+ uDux = u[1] .* setup.grid.Δu[1] .* setup.grid.Δ[2]' .* d[1]
+ uDuy = u[2] .* setup.grid.Δ[1] .* setup.grid.Δu[2]' .* d[2]
+ sum(uDux[Iu[1]]) + sum(uDuy[Iu[2]])
+ elseif D == 3
+ uDux =
+ u[1] .* setup.grid.Δu[1] .* reshape(setup.grid.Δ[2], 1, :) .*
+ reshape(setup.grid.Δ[3], 1, 1, :) .* d[1]
+ uDuy =
+ u[2] .* setup.grid.Δ[1] .* reshape(setup.grid.Δu[2], 1, :) .*
+ reshape(setup.grid.Δ[3], 1, 1, :) .* d[2]
+ uDuz =
+ u[3] .* setup.grid.Δ[1] .* reshape(setup.grid.Δ[2], 1, :) .*
+ reshape(setup.grid.Δu[3], 1, 1, :) .* d[3]
+ sum(uDux[Iu[1]]) + sum(uDuy[Iu[2]]) + sum(uDuz[Iu[3]])
+ end
+ @test d isa Tuple
+ @test d[1] isa Array{T}
+ @test uDu ≤ 0 # Check negativity (dissipation)
+ end
+
+ @testset "Convection-Diffusion" begin
+ cd = convectiondiffusion!(zero.(u), u, setup)
+ c_and_d = zero.(u)
+ convection!(c_and_d, u, setup)
+ diffusion!(c_and_d, u, setup)
+ @test all(cd .≈ c_and_d)
+ end
+
+ @testset "Momentum" begin
+ m = momentum(u, nothing, T(1), setup)
+ @test m isa Tuple
+ @test m[1] isa Array{T}
+ @test all(all.(!isnan, m))
+ end
+
+ @testset "Other fields" begin
+ p = randn!(similar(u[1]))
+ ω = vorticity(u, setup)
+ D == 2 && @test ω isa Array{T}
+ D == 3 && @test ω isa Tuple
+ @test smagorinsky_closure(setup)(u, 0.1) isa Tuple
+ @test tensorbasis(u, setup) isa Tuple
+ @test interpolate_u_p(u, setup) isa Tuple
+ D == 2 && @test interpolate_ω_p(ω, setup) isa Array{T}
+ D == 3 && @test interpolate_ω_p(ω, setup) isa Tuple
+ @test Dfield(p, setup) isa Array{T}
+ @test Qfield(u, setup) isa Array{T}
+ D == 2 && @test_throws AssertionError eig2field(u, setup)
+ D == 3 && @test eig2field(u, setup) isa Array{T} broken = D == 3
+ @test kinetic_energy(u, setup) isa Array{T}
+ @test total_kinetic_energy(u, setup) isa T
+ end
+end
+
+testops(IncompressibleNavierStokes.Dimension(2))
+testops(IncompressibleNavierStokes.Dimension(3))
diff --git a/test/postprocess2D.jl b/test/postprocess2D.jl
index 2881196f7..610fd4b10 100644
--- a/test/postprocess2D.jl
+++ b/test/postprocess2D.jl
@@ -7,9 +7,9 @@
setup = Setup(x, y)
- @test plot_grid(x, y) isa Makie.FigureAxisPlot
+ @test plotgrid(x, y) isa Makie.FigureAxisPlot
- pressure_solver = SpectralPressureSolver(setup)
+ psolver = psolver_spectral(setup)
t_start, t_end = tlims = (0.0, 1.0)
@@ -17,18 +17,18 @@
initial_velocity_u(x, y) = -sin(x)cos(y)
initial_velocity_v(x, y) = cos(x)sin(y)
initial_pressure(x, y) = 1 / 4 * (cos(2x) + cos(2y))
- V₀, p₀ = create_initial_conditions(
+ V = create_initial_conditions(
setup,
initial_velocity_u,
initial_velocity_v,
t_start;
initial_pressure,
- pressure_solver,
+ psolver,
)
# Iteration processors
processors = (
- field_plotter(setup; nupdate = 1, displayfig = false),
+ realtimeplotter(; setup, nupdate = 1, displayfig = false),
vtk_writer(setup; nupdate = 5, dir = "output", filename = "solution2D"),
animator(
setup,
@@ -36,12 +36,11 @@
nupdate = 10,
plotter = field_plotter(setup; displayfig = false),
),
- step_logger(),
+ timelogger(),
)
# Solve unsteady problem
- V, p, outputs =
- solve_unsteady(setup, V₀, p₀, tlims; Δt = 0.01, processors, pressure_solver)
+ state, outputs = solve_unsteady(setup, V₀, tlims; Δt = 0.01, processors, psolver)
@testset "VTK files" begin
@info "Testing 2D processors: VTK files"
diff --git a/test/postprocess3D.jl b/test/postprocess3D.jl
index 16d2fbe0f..eae19a86e 100644
--- a/test/postprocess3D.jl
+++ b/test/postprocess3D.jl
@@ -10,9 +10,9 @@
z = LinRange(lims..., n)
setup = Setup(x, y, z)
- @test plot_grid(x, y, z) isa Makie.Figure
+ @test plotgrid(x, y, z) isa Makie.Figure
- pressure_solver = SpectralPressureSolver(setup)
+ psolver = default_psolver(setup)
t_start, t_end = tlims = (T(0), T(1))
@@ -20,19 +20,19 @@
initial_velocity_v(x, y, z) = -cos(x)sin(y)cos(z)
initial_velocity_w(x, y, z) = zero(x)
initial_pressure(x, y, z) = 1 // 4 * (cos(2x) + cos(2y) + cos(2z))
- V₀, p₀ = create_initial_conditions(
+ V = create_initial_conditions(
setup,
initial_velocity_u,
initial_velocity_v,
initial_velocity_w,
t_start;
initial_pressure,
- pressure_solver,
+ psolver,
)
# Iteration processors
processors = (
- field_plotter(setup; nupdate = 5, displayfig = false),
+ realtimeplotter(; setup, nupdate = 5, displayfig = false),
vtk_writer(setup; nupdate = 5, dir = "output", filename = "solution3D"),
animator(
setup,
@@ -40,27 +40,17 @@
nupdate = 10,
plotter = field_plotter(setup; displayfig = false),
),
- step_logger(),
+ timelogger(),
)
# Solve unsteady problem
- # FIXME: using too many processors results in endless compilation
- V, p, outputs = solve_unsteady(
- setup,
- V₀,
- p₀,
- tlims;
- Δt = T(0.01),
- processors,
- pressure_solver,
- inplace = true,
- )
+ state, outputs = solve_unsteady(setup, u₀, tlims; Δt = T(0.01), processors, psolver)
@testset "VTK files" begin
@info "Testing 3D processors: VTK files"
@test isfile("output/solution3D.pvd")
@test isfile("output/solution3D_t=0p0.vti")
- save_vtk(setup, V, p, t_end, "output/field3D")
+ save_vtk(setup, state.u, state.v, t_end, "output/field3D")
@test isfile("output/field3D.vti")
end
diff --git a/test/pressure_solvers.jl b/test/pressure_solvers.jl
deleted file mode 100644
index 88988caf9..000000000
--- a/test/pressure_solvers.jl
+++ /dev/null
@@ -1,31 +0,0 @@
-@testset "Pressure solvers" begin
- @info "Testing pressure solvers"
- n = 20
- x = LinRange(0, 2π, n)
- y = LinRange(0, 2π, n)
- setup = Setup(x, y)
- (; A) = setup.operators
-
- # Pressure solvers
- direct = DirectPressureSolver(setup)
- cg = CGPressureSolver(setup)
- spectral = SpectralPressureSolver(setup)
-
- initial_pressure(x, y) = 1 / 4 * (cos(2x) + cos(2y))
- p_exact = reshape(initial_pressure.(setup.grid.xpp, setup.grid.ypp), :)
- f = A * p_exact
-
- p_direct = pressure_poisson(direct, f)
- p_cg = pressure_poisson(cg, f)
- p_spectral = pressure_poisson(spectral, f)
-
- # Test that in-place and out-of-place versions give same result
- @test p_direct ≈ pressure_poisson!(direct, zero(p_exact), f)
- @test p_cg ≈ pressure_poisson!(cg, zero(p_exact), f)
- @test p_spectral ≈ pressure_poisson!(spectral, zero(p_exact), f)
-
- # Test that solvers compute the exact pressure
- @test_broken p_direct ≈ p_exact # `A` is really badly conditioned
- @test p_cg ≈ p_exact
- @test p_spectral ≈ p_exact
-end
diff --git a/test/psolvers.jl b/test/psolvers.jl
new file mode 100644
index 000000000..2391667de
--- /dev/null
+++ b/test/psolvers.jl
@@ -0,0 +1,36 @@
+@testset "Pressure solvers" begin
+ @info "Testing pressure solvers"
+ n = 32
+ x = LinRange(0, 2π, n + 1)
+ y = LinRange(0, 2π, n + 1)
+ Re = 1e3
+ setup = Setup(x, y; Re)
+ (; xp) = setup.grid
+ D = 2
+
+ initial_pressure(x, y) = 1 / 4 * (cos(2x) + cos(2y))
+ p_exact =
+ initial_pressure.(
+ ntuple(α -> reshape(xp[α], ntuple(Returns(1), α - 1)..., :), D)...,
+ )
+ IncompressibleNavierStokes.apply_bc_p!(p_exact, 0.0, setup)
+ lap = IncompressibleNavierStokes.laplacian(p_exact, setup)
+
+ # Pressure solvers
+ direct = psolver_direct(setup)
+ cg = psolver_cg(setup)
+ spectral = psolver_spectral(setup)
+ spectral_lowmemory = psolver_spectral_lowmemory(setup)
+
+ get_p(psolver) = IncompressibleNavierStokes.apply_bc_p(
+ IncompressibleNavierStokes.poisson(psolver, lap),
+ 0.0,
+ setup,
+ )
+
+ # Test that solvers compute the exact pressure
+ @test get_p(direct) ≈ p_exact
+ @test get_p(cg) ≈ p_exact
+ @test get_p(spectral) ≈ p_exact
+ @test get_p(spectral_lowmemory) ≈ p_exact
+end
diff --git a/test/runtests.jl b/test/runtests.jl
index ebedef16b..eb0205df7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,32 +1,56 @@
-# LSP indexing solution
-# https://github.com/julia-vscode/julia-vscode/issues/800#issuecomment-650085983
-if isdefined(@__MODULE__, :LanguageServer)
- include("../src/IncompressibleNavierStokes.jl")
- using .IncompressibleNavierStokes
-end
+push!(LOAD_PATH, joinpath(@__DIR__, ".."))
using Aqua
using CairoMakie
+using ChainRulesCore
+using ChainRulesTestUtils
using IncompressibleNavierStokes
+using IncompressibleNavierStokes:
+ apply_bc_p,
+ apply_bc_u,
+ bodyforce,
+ convection,
+ convection!,
+ convectiondiffusion!,
+ diffusion,
+ diffusion!,
+ divergence,
+ eig2field,
+ kinetic_energy,
+ interpolate_u_p,
+ interpolate_ω_p,
+ laplacian,
+ laplacian_mat,
+ momentum,
+ poisson,
+ pressuregradient,
+ smagorinsky_closure,
+ tensorbasis,
+ total_kinetic_energy,
+ vorticity,
+ Dfield,
+ Qfield
+
using LinearAlgebra
+using Random
+using SparseArrays
using Statistics
using Test
@testset "IncompressibleNavierStokes" begin
include("grid.jl")
- include("pressure_solvers.jl")
- include("models.jl")
- include("solvers.jl")
- include("simulation2D.jl")
- include("simulation3D.jl")
- include("postprocess2D.jl")
- include("postprocess3D.jl")
+ include("psolvers.jl")
+ include("operators.jl")
+ include("chainrules.jl")
+ # include("models.jl")
+ # include("solvers.jl")
+ # include("simulation2D.jl")
+ # include("simulation3D.jl")
+ # include("postprocess2D.jl")
+ # include("postprocess3D.jl")
@testset "Aqua" begin
@info "Testing code with Aqua"
- Aqua.test_all(
- IncompressibleNavierStokes;
- ambiguities = false,
- )
+ Aqua.test_all(IncompressibleNavierStokes; ambiguities = false)
end
end
diff --git a/test/simulation2D.jl b/test/simulation2D.jl
index f12c865ac..a915f4f02 100644
--- a/test/simulation2D.jl
+++ b/test/simulation2D.jl
@@ -20,7 +20,7 @@
initial_velocity_u(x, y) = 0.0
initial_velocity_v(x, y) = 0.0
initial_pressure(x, y) = 0.0
- V₀, p₀ = create_initial_conditions(
+ V = create_initial_conditions(
setup,
initial_velocity_u,
initial_velocity_v,
@@ -40,16 +40,15 @@
end
# Iteration processors
- processors = (step_logger(),)
+ processors = (timelogger(),)
@testset "Unsteady problem" begin
- V, p, outputs = solve_unsteady(setup, V₀, p₀, tlims; Δt = 0.01, processors)
+ (; u, t), outputs = solve_unsteady(setup, V₀, tlims; Δt = 0.01, processors)
# Check that solution did not explode
- @test all(!isnan, V)
- @test all(!isnan, p)
+ @test all(!isnan, u)
# Check that the average velocity is smaller than the lid velocity
- @test sum(abs, V) / length(V) < lid_vel
+ @test sum(abs, u) / length(u) < lid_vel
end
end
diff --git a/test/simulation3D.jl b/test/simulation3D.jl
index f1851c65e..7397fb4c2 100644
--- a/test/simulation3D.jl
+++ b/test/simulation3D.jl
@@ -18,7 +18,7 @@
initial_velocity_v(x, y, z) = 0.0
initial_velocity_w(x, y, z) = 0.0
initial_pressure(x, y, z) = 0.0
- V₀, p₀ = create_initial_conditions(
+ V = create_initial_conditions(
setup,
initial_velocity_u,
initial_velocity_v,
@@ -28,27 +28,27 @@
)
@testset "Steady state problem" begin
- V, p = solve_steady_state(setup, V₀, p₀)
+ u, p = solve_steady_state(setup, V₀, p₀)
# Check that solution did not explode
- @test all(!isnan, V)
+ @test all(!isnan, u)
@test all(!isnan, p)
# Check that the average velocity is smaller than the lid velocity
- @test sum(abs, V) / length(V) < norm(lid_vel)
+ @test sum(abs, u) / length(u) < norm(lid_vel)
end
# Iteration processors
- processors = (step_logger(),)
+ processors = (timelogger(),)
@testset "Unsteady problem" begin
- V, p, outputs = solve_unsteady(setup, V₀, p₀, tlims; Δt = 0.01, processors)
+ (; u, t), outputs = solve_unsteady(setup, V₀, tlims; Δt = 0.01, processors)
# Check that solution did not explode
- @test all(!isnan, V)
+ @test all(!isnan, u)
@test all(!isnan, p)
# Check that the average velocity is smaller than the lid velocity
- @test sum(abs, V) / length(V) < norm(lid_vel)
+ @test sum(abs, u) / length(u) < norm(lid_vel)
end
end
diff --git a/test/solvers.jl b/test/solvers.jl
index d5cbd37b5..91f43ca61 100644
--- a/test/solvers.jl
+++ b/test/solvers.jl
@@ -1,27 +1,26 @@
@testset "Solvers" begin
T = Float64
Re = 500.0
- viscosity_model = LaminarModel(; Re)
n = 50
x = LinRange(0, 2π, n + 1)
y = LinRange(0, 2π, n + 1)
- setup = Setup(x, y; viscosity_model)
+ setup = Setup(x, y; Re)
- pressure_solver = SpectralPressureSolver(setup)
+ psolver = default_psolver(setup)
t_start, t_end = tlims = (0.0, 5.0)
initial_velocity_u(x, y) = cos(x)sin(y)
initial_velocity_v(x, y) = -sin(x)cos(y)
initial_pressure(x, y) = -1 / 4 * (cos(2x) + cos(2y))
- V₀, p₀ = create_initial_conditions(
+ V = create_initial_conditions(
setup,
initial_velocity_u,
initial_velocity_v,
t_start;
initial_pressure,
- pressure_solver,
+ psolver,
)
@testset "Steady state" begin
@@ -45,28 +44,14 @@
@testset "Unsteady solvers" begin
@testset "Explicit Runge Kutta" begin
@info "Testing explicit Runge-Kutta, out-of-place version"
- V, p, outputs = solve_unsteady(
- setup,
- V₀,
- p₀,
- tlims;
- Δt = 0.01,
- pressure_solver,
- inplace = false,
- )
- @test norm(V - V_exact) / norm(V_exact) < 1e-4
+ state, outputs =
+ solve_unsteady(setup, V₀, tlims; Δt = 0.01, psolver, inplace = false)
+ @test norm(state.u - u_exact) / norm(u_exact) < 1e-4
@info "Testing explicit Runge-Kutta, in-place version"
- Vip, pip, outputsip = solve_unsteady(
- setup,
- V₀,
- p₀,
- tlims;
- Δt = 0.01,
- pressure_solver,
- inplace = true,
- )
- @test Vip ≈ V
- @test pip ≈ p
+ stateip, outputsip =
+ solve_unsteady(setup, V₀, tlims; Δt = 0.01, psolver, inplace = true)
+ @test stateip.u ≈ state.u
+ @test stateip.p ≈ state.p
end
@testset "Implicit Runge Kutta" begin
@@ -74,82 +59,76 @@
@test_broken solve_unsteady(
setup,
V₀,
- p₀,
tlims;
method = RIA2(),
Δt = 0.01,
- pressure_solver,
+ psolver,
inplace = false,
) isa Tuple
@info "Testing implicit Runge-Kutta, in-place version"
- V, p, outputs = solve_unsteady(
+ (; u, t), outputs = solve_unsteady(
setup,
V₀,
- p₀,
tlims;
method = RIA2(),
Δt = 0.01,
- pressure_solver,
+ psolver,
inplace = true,
- processors = (step_logger(),),
+ processors = (timelogger(),),
)
- @test_broken norm(V - V_exact) / norm(V_exact) < 1e-3
+ @test_broken norm(u - u_exact) / norm(u_exact) < 1e-3
end
@testset "One-leg beta method" begin
@info "Testing one-leg beta method, out-of-place version"
- V, p, outputs = solve_unsteady(
+ state, outputs = solve_unsteady(
setup,
V₀,
- p₀,
tlims;
method = OneLegMethod(T),
Δt = 0.01,
- pressure_solver,
+ psolver,
inplace = false,
)
- @test norm(V - V_exact) / norm(V_exact) < 1e-4
+ @test norm(state.u - u_exact) / norm(u_exact) < 1e-4
@info "Testing one-leg beta method, in-place version"
- Vip, pip, outputsip = solve_unsteady(
+ stateip, outputsip = solve_unsteady(
setup,
V₀,
- p₀,
tlims;
method = OneLegMethod(T),
Δt = 0.01,
- pressure_solver,
+ psolver,
inplace = true,
)
- @test Vip ≈ V
- @test pip ≈ p
+ @test stateip.u ≈ state.u
+ @test stateip.p ≈ state.p
end
@testset "Adams-Bashforth Crank-Nicolson" begin
@info "Testing Adams-Bashforth Crank-Nicolson method, out-of-place version"
- V, p, outputs = solve_unsteady(
+ state, outputs = solve_unsteady(
setup,
V₀,
- p₀,
tlims;
method = AdamsBashforthCrankNicolsonMethod(T),
Δt = 0.01,
- pressure_solver,
+ psolver,
inplace = false,
)
- @test norm(V - V_exact) / norm(V_exact) < 1e-4
+ @test norm(state.u - u_exact) / norm(u_exact) < 1e-4
@info "Testing Adams-Bashforth Crank-Nicolson method, in-place version"
- Vip, pip, outputs = solve_unsteady(
+ stateip, outputs = solve_unsteady(
setup,
V₀,
- p₀,
tlims;
method = AdamsBashforthCrankNicolsonMethod(T),
Δt = 0.01,
- pressure_solver,
+ psolver,
inplace = true,
)
- @test Vip ≈ V
- @test pip ≈ p
+ @test stateip.u ≈ state.u
+ @test stateip.p ≈ state.p
end
end
end