Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 178 additions & 0 deletions demos/extruded_shallow_water/extruded_shallow_water.py.rst
Original file line number Diff line number Diff line change
@@ -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>`.
80 changes: 0 additions & 80 deletions demos/extruded_shallow_water/test_extrusion_lsw.py

This file was deleted.

1 change: 1 addition & 0 deletions docs/source/intro_tut.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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>
1 change: 1 addition & 0 deletions tests/firedrake/demos/test_demos_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"), []),
Expand Down