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..d9f752d7d7 --- /dev/null +++ b/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 `. \ 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) 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. 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"), []),