From 4e2880f1edf64850e29fdb2b4c9eab2b35818507 Mon Sep 17 00:00:00 2001 From: ZiyaSelim Date: Sat, 31 Jan 2026 15:53:43 +0100 Subject: [PATCH 1/8] Change python script to rst demo --- .../extruded_shallow_water.py.rst | 124 ++++++++++++++++++ .../test_extrusion_lsw.py | 80 ----------- 2 files changed, 124 insertions(+), 80 deletions(-) create mode 100644 demos/extruded_shallow_water/extruded_shallow_water.py.rst delete mode 100644 demos/extruded_shallow_water/test_extrusion_lsw.py diff --git a/demos/extruded_shallow_water/extruded_shallow_water.py.rst b/demos/extruded_shallow_water/extruded_shallow_water.py.rst new file mode 100644 index 0000000000..dfa5ab1e89 --- /dev/null +++ b/demos/extruded_shallow_water/extruded_shallow_water.py.rst @@ -0,0 +1,124 @@ +Linear Shallow Water Equations on an Extruded Mesh +================================================== + +This demo solves the linear shallow water equations on an extruded mesh +using a Strang timestepping scheme. + +The equations are solved in a domain $\Omega$ utilizing a 2D base mesh +that is extruded vertically to form a 3D volume. + +As usual, we start by importing Firedrake: :: + + from firedrake import * + +Mesh Generation +---------------- + +We use an *extruded* mesh, where the base mesh is a $2^5 \times 2^5$ unit square +with 5 evenly-spaced vertical layers. This results in a 3D volume composed of +prisms. :: + + power = 5 + m = UnitSquareMesh(2 ** power, 2 ** power) + layers = 5 + + mesh = ExtrudedMesh(m, layers, layer_height=0.25) + +Function Spaces +---------------- +For the velocity field, we use an $H(\mathrm{div})$-conforming function space +constructed as the outer product of a 2D BDM space and a 1D DG space. This ensures +that the normal component of the velocity is continuous across element boundaries +in the horizontal directions, which is important for accurately capturing fluxes. :: + + horiz = FiniteElement("BDM", "triangle", 1) + vert = FiniteElement("DG", "interval", 0) + prod = HDiv(OuterProductElement(horiz, vert)) + W = FunctionSpace(mesh, prod) + +We also define a pressure space $X$ using piecewise constant discontinuous +Galerkin elements, and a plotting space $X_{\text{plot}}$ using continuous +Galerkin elements for better visualization. :: + + X = FunctionSpace(mesh, "DG", 0, vfamily="DG", vdegree=0) + Xplot = FunctionSpace(mesh, "CG", 1, vfamily="Lagrange", vdegree=1) + +Initial Conditions +------------------- + +We define our functions for velocity and pressure fields. The initial pressure field +is set to a prescribed sine function to create a wave-like disturbance. :: + + # Define starting field + u_0 = Function(W) + u_h = Function(W) + u_1 = Function(W) + p_0 = Function(X) + p_1 = Function(X) + p_plot = Function(Xplot) + x, y = SpatialCoordinate(m) + p_0.interpolate(sin(4*pi*x)*sin(2*pi*x)) + + T = 0.5 + t = 0 + dt = 0.0025 + + file = VTKFile("lsw3d.pvd") + +Before starting the time-stepping loop, we project the initial pressure field +into the plotting space for visualization. :: + + p_trial = TrialFunction(Xplot) + p_test = TestFunction(Xplot) + solve(p_trial * p_test * dx == p_0 * p_test * dx, p_plot) + file << p_plot, t + + E_0 = assemble(0.5 * p_0 * p_0 * dx + 0.5 * dot(u_0, u_0) * dx) + +Time-Stepping Loop +-------------------- + +We evolve the system in time using a Strang splitting method. In each time step, +we perform a half-step update of the velocity field, a full-step update of the pressure field, +and then another half-step update of the velocity field. This approach helps to +maintain stability and accuracy. :: + + while t < T: + u = TrialFunction(W) + w = TestFunction(W) + a_1 = dot(w, u) * dx + L_1 = dot(w, u_0) * dx + 0.5 * dt * div(w) * p_0 * dx + solve(a_1 == L_1, u_h) + + p = TrialFunction(X) + phi = TestFunction(X) + a_2 = phi * p * dx + L_2 = phi * p_0 * dx - dt * phi * div(u_h) * dx + solve(a_2 == L_2, p_1) + + u = TrialFunction(W) + w = TestFunction(W) + a_3 = dot(w, u) * dx + L_3 = dot(w, u_h) * dx + 0.5 * dt * div(w) * p_1 * dx + solve(a_3 == L_3, u_1) + + u_0.assign(u_1) + p_0.assign(p_1) + t += dt + + # project into P1 x P1 for plotting + p_trial = TrialFunction(Xplot) + p_test = TestFunction(Xplot) + solve(p_trial * p_test * dx == p_0 * p_test * dx, p_plot) + file << p_plot, t + print(t) + +Energy Calculation +------------------ + +Finally, we compute and print the total energy of the system at the end of the simulation +and compare it to the initial energy to assess conservation properties. :: + + E_1 = assemble(0.5 * p_0 * p_0 * dx + 0.5 * dot(u_0, u_0) * dx) + print('Initial energy', E_0) + print('Final energy', E_1) \ No newline at end of file diff --git a/demos/extruded_shallow_water/test_extrusion_lsw.py b/demos/extruded_shallow_water/test_extrusion_lsw.py deleted file mode 100644 index be1719dbc1..0000000000 --- a/demos/extruded_shallow_water/test_extrusion_lsw.py +++ /dev/null @@ -1,80 +0,0 @@ -# FIXME: document properly -"""Demo of Linear Shallow Water, with Strang timestepping and silly BCs, but -a sin(x)*sin(y) solution that doesn't care about the silly BCs""" - -from firedrake import * - - -power = 5 -# Create mesh and define function space -m = UnitSquareMesh(2 ** power, 2 ** power) -layers = 5 - -# Populate the coordinates of the extruded mesh by providing the -# coordinates as a field. - -mesh = ExtrudedMesh(m, layers, layer_height=0.25) - -horiz = FiniteElement("BDM", "triangle", 1) -vert = FiniteElement("DG", "interval", 0) -prod = HDiv(OuterProductElement(horiz, vert)) -W = FunctionSpace(mesh, prod) - -X = FunctionSpace(mesh, "DG", 0, vfamily="DG", vdegree=0) -Xplot = FunctionSpace(mesh, "CG", 1, vfamily="Lagrange", vdegree=1) - -# Define starting field -u_0 = Function(W) -u_h = Function(W) -u_1 = Function(W) -p_0 = Function(X) -p_1 = Function(X) -p_plot = Function(Xplot) -x, y = SpatialCoordinate(m) -p_0.interpolate(sin(4*pi*x)*sin(2*pi*x)) - -T = 0.5 -t = 0 -dt = 0.0025 - -file = VTKFile("lsw3d.pvd") -p_trial = TrialFunction(Xplot) -p_test = TestFunction(Xplot) -solve(p_trial * p_test * dx == p_0 * p_test * dx, p_plot) -file << p_plot, t - -E_0 = assemble(0.5 * p_0 * p_0 * dx + 0.5 * dot(u_0, u_0) * dx) - -while t < T: - u = TrialFunction(W) - w = TestFunction(W) - a_1 = dot(w, u) * dx - L_1 = dot(w, u_0) * dx + 0.5 * dt * div(w) * p_0 * dx - solve(a_1 == L_1, u_h) - - p = TrialFunction(X) - phi = TestFunction(X) - a_2 = phi * p * dx - L_2 = phi * p_0 * dx - dt * phi * div(u_h) * dx - solve(a_2 == L_2, p_1) - - u = TrialFunction(W) - w = TestFunction(W) - a_3 = dot(w, u) * dx - L_3 = dot(w, u_h) * dx + 0.5 * dt * div(w) * p_1 * dx - solve(a_3 == L_3, u_1) - - u_0.assign(u_1) - p_0.assign(p_1) - t += dt - - # project into P1 x P1 for plotting - p_trial = TrialFunction(Xplot) - p_test = TestFunction(Xplot) - solve(p_trial * p_test * dx == p_0 * p_test * dx, p_plot) - file << p_plot, t - print(t) - -E_1 = assemble(0.5 * p_0 * p_0 * dx + 0.5 * dot(u_0, u_0) * dx) -print('Initial energy', E_0) -print('Final energy', E_1) From 115de2ff34c74330ed5346042f1064ff89382491 Mon Sep 17 00:00:00 2001 From: ZiyaSelim Date: Sat, 31 Jan 2026 15:54:02 +0100 Subject: [PATCH 2/8] Add new rst demo to test file for demos --- tests/firedrake/demos/test_demos_run.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/firedrake/demos/test_demos_run.py b/tests/firedrake/demos/test_demos_run.py index dbb0b938be..bd46952e55 100644 --- a/tests/firedrake/demos/test_demos_run.py +++ b/tests/firedrake/demos/test_demos_run.py @@ -26,6 +26,7 @@ Demo(("DG_advection", "DG_advection"), ["matplotlib"]), Demo(("eigenvalues_QG_basinmodes", "qgbasinmodes"), ["matplotlib", "slepc", "vtk"]), Demo(("extruded_continuity", "extruded_continuity"), []), + Demo(("extruded_shallow_water", "extruded_shallow_water"), []), Demo(("helmholtz", "helmholtz"), ["vtk"]), Demo(("higher_order_mass_lumping", "higher_order_mass_lumping"), ["vtk"]), Demo(("immersed_fem", "immersed_fem"), []), From 0f285b85fbf564dcab0c0b149050da67c807d093 Mon Sep 17 00:00:00 2001 From: ZiyaSelim Date: Tue, 3 Feb 2026 19:43:34 +0100 Subject: [PATCH 3/8] Add python script to rstdemo --- .../extruded_shallow_water/extruded_shallow_water.py.rst | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/demos/extruded_shallow_water/extruded_shallow_water.py.rst b/demos/extruded_shallow_water/extruded_shallow_water.py.rst index dfa5ab1e89..3a0775eb1e 100644 --- a/demos/extruded_shallow_water/extruded_shallow_water.py.rst +++ b/demos/extruded_shallow_water/extruded_shallow_water.py.rst @@ -1,5 +1,5 @@ -Linear Shallow Water Equations on an Extruded Mesh -================================================== +Linear Shallow Water Equations on an Extruded Mesh using a Strang timestepping scheme +===================================================================================== This demo solves the linear shallow water equations on an extruded mesh using a Strang timestepping scheme. @@ -121,4 +121,7 @@ and compare it to the initial energy to assess conservation properties. :: E_1 = assemble(0.5 * p_0 * p_0 * dx + 0.5 * dot(u_0, u_0) * dx) print('Initial energy', E_0) - print('Final energy', E_1) \ No newline at end of file + print('Final energy', E_1) + +This demo can be found as a script in +:demo:`extruded_shallow_water.py `. \ No newline at end of file From bdf7927439b8376b1d5708acf0753f72503dfe1b Mon Sep 17 00:00:00 2001 From: ZiyaSelim Date: Tue, 3 Feb 2026 19:43:58 +0100 Subject: [PATCH 4/8] Add new demo to intro --- docs/source/intro_tut.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/intro_tut.rst b/docs/source/intro_tut.rst index 6f31f420b1..93b3d3afeb 100644 --- a/docs/source/intro_tut.rst +++ b/docs/source/intro_tut.rst @@ -14,5 +14,6 @@ Introductory Tutorials A mixed formulation of the Poisson equation. A time-dependent DG advection equation using upwinding. An extruded mesh example, using a steady-state continuity equation. + A linear shallow water equations example using a Strang timestepping scheme. A linear wave equation with optional mass lumping. Creating Firedrake-compatible meshes in Gmsh. From 931e95eb474592a66d8b6c1ef4e7d98826e5e17e Mon Sep 17 00:00:00 2001 From: ZiyaSelim Date: Wed, 4 Feb 2026 20:11:51 +0100 Subject: [PATCH 5/8] Use inline math text instead of $$ --- .../extruded_shallow_water.py.rst | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/demos/extruded_shallow_water/extruded_shallow_water.py.rst b/demos/extruded_shallow_water/extruded_shallow_water.py.rst index 3a0775eb1e..b33f820a3f 100644 --- a/demos/extruded_shallow_water/extruded_shallow_water.py.rst +++ b/demos/extruded_shallow_water/extruded_shallow_water.py.rst @@ -4,7 +4,7 @@ Linear Shallow Water Equations on an Extruded Mesh using a Strang timestepping s This demo solves the linear shallow water equations on an extruded mesh using a Strang timestepping scheme. -The equations are solved in a domain $\Omega$ utilizing a 2D base mesh +The equations are solved in a domain :math:`\Omega` utilizing a 2D base mesh that is extruded vertically to form a 3D volume. As usual, we start by importing Firedrake: :: @@ -14,19 +14,19 @@ As usual, we start by importing Firedrake: :: Mesh Generation ---------------- -We use an *extruded* mesh, where the base mesh is a $2^5 \times 2^5$ unit square +We use an *extruded* mesh, where the base mesh is a :math:`2^5 \times 2^5` unit square with 5 evenly-spaced vertical layers. This results in a 3D volume composed of prisms. :: - power = 5 - m = UnitSquareMesh(2 ** power, 2 ** power) - layers = 5 + power = 5 + m = UnitSquareMesh(2 ** power, 2 ** power) + layers = 5 - mesh = ExtrudedMesh(m, layers, layer_height=0.25) + mesh = ExtrudedMesh(m, layers, layer_height=0.25) Function Spaces ---------------- -For the velocity field, we use an $H(\mathrm{div})$-conforming function space +For the velocity field, we use an :math:`H(\mathrm{div})`-conforming function space constructed as the outer product of a 2D BDM space and a 1D DG space. This ensures that the normal component of the velocity is continuous across element boundaries in the horizontal directions, which is important for accurately capturing fluxes. :: @@ -36,8 +36,8 @@ in the horizontal directions, which is important for accurately capturing fluxes prod = HDiv(OuterProductElement(horiz, vert)) W = FunctionSpace(mesh, prod) -We also define a pressure space $X$ using piecewise constant discontinuous -Galerkin elements, and a plotting space $X_{\text{plot}}$ using continuous +We also define a pressure space :math:`X` using piecewise constant discontinuous +Galerkin elements, and a plotting space :math:`X_{\text{plot}}` using continuous Galerkin elements for better visualization. :: X = FunctionSpace(mesh, "DG", 0, vfamily="DG", vdegree=0) From 93cb02402849950d9cf806dcd2b323d56d7f0a83 Mon Sep 17 00:00:00 2001 From: ZiyaSelim Date: Thu, 5 Feb 2026 21:15:07 +0100 Subject: [PATCH 6/8] Add math block explanations to single steps of the time stepping loop --- .../extruded_shallow_water.py.rst | 53 ++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/demos/extruded_shallow_water/extruded_shallow_water.py.rst b/demos/extruded_shallow_water/extruded_shallow_water.py.rst index b33f820a3f..62a5728cc7 100644 --- a/demos/extruded_shallow_water/extruded_shallow_water.py.rst +++ b/demos/extruded_shallow_water/extruded_shallow_water.py.rst @@ -81,7 +81,58 @@ Time-Stepping Loop We evolve the system in time using a Strang splitting method. In each time step, we perform a half-step update of the velocity field, a full-step update of the pressure field, and then another half-step update of the velocity field. This approach helps to -maintain stability and accuracy. :: +maintain stability and accuracy. + +**Step 1: Velocity Half-Step** +First, we solve for an intermediate velocity :math:`u_h` using the pressure +from the start of the step :math:`p_0`. Mathematically, we find :math:`u_h \in W` such that: + +.. math:: + + \int_{\Omega} w \cdot u_h \, dx = \int_{\Omega} w \cdot u_0 \, dx + \frac{\Delta t}{2} \int_{\Omega} (\nabla \cdot w) p_0 \, dx \quad \forall w \in W + +.. code-block:: python + + a_1 = dot(w, u) * dx + L_1 = dot(w, u_0) * dx + 0.5 * dt * div(w) * p_0 * dx + solve(a_1 == L_1, u_h) + +**Step 2: Pressure Full-Step** +Next, we update the pressure to :math:`p_1` using the divergence of the +intermediate velocity :math:`u_h`. We find :math:`p_1 \in X` such that: + +.. math:: + + \int_{\Omega} \phi \, p_1 \, dx = \int_{\Omega} \phi \, p_0 \, dx - \Delta t \int_{\Omega} \phi (\nabla \cdot u_h) \, dx \quad \forall \phi \in X + +.. code-block:: python + + a_2 = phi * p * dx + L_2 = phi * p_0 * dx - dt * phi * div(u_h) * dx + solve(a_2 == L_2, p_1) + +**Step 3: Velocity Final Half-Step** +Finally, we compute the velocity at the end of the time step :math:`u_1` using +the updated pressure :math:`p_1`. We find :math:`u_1 \in W` such that: + +.. math:: + + \int_{\Omega} w \cdot u_1 \, dx = \int_{\Omega} w \cdot u_h \, dx + \frac{\Delta t}{2} \int_{\Omega} (\nabla \cdot w) p_1 \, dx \quad \forall w \in W + +.. code-block:: python + + a_3 = dot(w, u) * dx + L_3 = dot(w, u_h) * dx + 0.5 * dt * div(w) * p_1 * dx + solve(a_3 == L_3, u_1) + +Implementation in the Simulation Loop +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Here follows the complete implementation of the time-stepping loop, including the updates of the velocity and pressure fields, +as well as the projection of the pressure field for visualization at each time step. +We also print the current simulation time at each step for tracking progress. + +.. code-block:: python while t < T: u = TrialFunction(W) From 50795513b021cda82b6a3c1de9613312c2fbfe97 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Fri, 6 Feb 2026 15:26:14 +0000 Subject: [PATCH 7/8] Update demos/extruded_shallow_water/extruded_shallow_water.py.rst Co-authored-by: Josh Hope-Collins --- demos/extruded_shallow_water/extruded_shallow_water.py.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/demos/extruded_shallow_water/extruded_shallow_water.py.rst b/demos/extruded_shallow_water/extruded_shallow_water.py.rst index 62a5728cc7..f38c01c367 100644 --- a/demos/extruded_shallow_water/extruded_shallow_water.py.rst +++ b/demos/extruded_shallow_water/extruded_shallow_water.py.rst @@ -33,7 +33,7 @@ in the horizontal directions, which is important for accurately capturing fluxes horiz = FiniteElement("BDM", "triangle", 1) vert = FiniteElement("DG", "interval", 0) - prod = HDiv(OuterProductElement(horiz, vert)) + prod = HDiv(TensorProductElement(horiz, vert)) W = FunctionSpace(mesh, prod) We also define a pressure space :math:`X` using piecewise constant discontinuous From 9a72e8bd856f11b81571d5cb3404ca134b8ce356 Mon Sep 17 00:00:00 2001 From: ZiyaSelim Date: Tue, 10 Feb 2026 15:42:21 +0100 Subject: [PATCH 8/8] Fix interpolation and file writing error in demo script --- demos/extruded_shallow_water/extruded_shallow_water.py.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/demos/extruded_shallow_water/extruded_shallow_water.py.rst b/demos/extruded_shallow_water/extruded_shallow_water.py.rst index f38c01c367..d9f752d7d7 100644 --- a/demos/extruded_shallow_water/extruded_shallow_water.py.rst +++ b/demos/extruded_shallow_water/extruded_shallow_water.py.rst @@ -56,7 +56,7 @@ is set to a prescribed sine function to create a wave-like disturbance. :: p_0 = Function(X) p_1 = Function(X) p_plot = Function(Xplot) - x, y = SpatialCoordinate(m) + x, y, z = SpatialCoordinate(mesh) p_0.interpolate(sin(4*pi*x)*sin(2*pi*x)) T = 0.5 @@ -71,7 +71,7 @@ into the plotting space for visualization. :: p_trial = TrialFunction(Xplot) p_test = TestFunction(Xplot) solve(p_trial * p_test * dx == p_0 * p_test * dx, p_plot) - file << p_plot, t + file.write(p_plot, time=t) E_0 = assemble(0.5 * p_0 * p_0 * dx + 0.5 * dot(u_0, u_0) * dx) @@ -161,7 +161,7 @@ We also print the current simulation time at each step for tracking progress. p_trial = TrialFunction(Xplot) p_test = TestFunction(Xplot) solve(p_trial * p_test * dx == p_0 * p_test * dx, p_plot) - file << p_plot, t + file.write(p_plot, time=t) print(t) Energy Calculation