Change test_extrusion_lsw python demo to rst demo (#4852)

---------

Co-authored-by: Connor Ward <c.ward20@imperial.ac.uk>
Co-authored-by: Josh Hope-Collins <jhc.jss@gmail.com>

Author: 3 people

demos/extruded_shallow_water/extruded_shallow_water.py.rst

@@ -0,0 +1,178 @@
+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.
+
+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: ::
+
+  from firedrake import *
+
+Mesh Generation
+----------------
+
+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
+
+  mesh = ExtrudedMesh(m, layers, layer_height=0.25)
+
+Function Spaces
+----------------
+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. ::
+
+  horiz = FiniteElement("BDM", "triangle", 1)
+  vert = FiniteElement("DG", "interval", 0)
+  prod = HDiv(TensorProductElement(horiz, vert))
+  W = FunctionSpace(mesh, prod)
+
+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)
+  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, z = SpatialCoordinate(mesh)
+  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.write(p_plot, time=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.
+
+**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)
+    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.write(p_plot, time=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)
+
+This demo can be found as a script in
+:demo:`extruded_shallow_water.py <extruded_shallow_water.py>`.

demos/extruded_shallow_water/test_extrusion_lsw.py

@@ -14,5 +14,6 @@ Introductory Tutorials
    A mixed formulation of the Poisson equation.<demos/poisson_mixed.py>
    A time-dependent DG advection equation using upwinding.<demos/DG_advection.py>
    An extruded mesh example, using a steady-state continuity equation.<demos/extruded_continuity.py>
+   A linear shallow water equations example using a Strang timestepping scheme.<demos/extruded_shallow_water.py>
    A linear wave equation with optional mass lumping.<demos/linear_wave_equation.py>
    Creating Firedrake-compatible meshes in Gmsh.<demos/immersed_fem.py>

docs/source/intro_tut.rst

@@ -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"), []),