Download this example.

2D advection-diffusion testcase

This example solves a simple shallow-water and tracer transport problem on a square domain. Two tracer components are transported by an imposed inflow and diffuse inside the domain.

The setup is mainly intended as a compact testcase for tracer advection, diffusion, and open-boundary forcing.

import numpy as np
import slim4
from slim4 import slim, slim2d, tools

First, we load the mesh. If the simulation is run in parallel, the mesh is partitioned automatically before being loaded.

g = 9.80616
meshfile = "square.msh"

if slim4.mpi.size > 1:
    partitioned_meshfile = slim2d.partition_mesh(meshfile)
    mesh = slim2d.Mesh(partitioned_meshfile)
else:
    mesh = slim2d.Mesh(meshfile)

We define the bathymetry and create the shallow-water and tracer problems. The tracer problem carries two components.

nf = 200
bathymetry = slim.DataStatic.create_constant(mesh, [6])
shallow_water = slim2d.ShallowWater(mesh, g, bathymetry)
tracer_problem = slim2d.Tracer(mesh, shallow_water, bathymetry, g, nf=nf)
solution = tracer_problem.new_vector()
print(solution.size, solution.size < 2**31)

We initialize the prognostic fields. The shallow-water state is prescribed in terms of free-surface elevation and transport, while both tracer components start from zero inside the domain.

transport_state = solution.get_sub(0).get()
tracer_state = solution.get_sub(1).get()

transport_state[0] = 0.0
transport_state[1] = 0.0
transport_state[2] = 2 * 6
tracer_state[0] = 0.0
tracer_state[1] = 0.0

solution.get_sub(0).set(transport_state)
solution.get_sub(1).set(tracer_state)

Next, we set the tracer diffusion and source term.

# horizontal_diffusivity = slim.DataStatic.create_constant(mesh, 100)
tracer_source = slim.DataStatic.create_constant(mesh, 0.0)

# tracer_problem.set_diffusion(horizontal_diffusivity)
tracer_problem.set_source(tracer_source)

We define the open-boundary data. The same tracer concentration, free-surface elevation, and transport are imposed on both open boundaries.

open_boundary_tracer = slim.DataStatic.create_constant(mesh, 1+np.arange(nf))
open_boundary_eta = slim.DataStatic.create_constant(mesh, 0)
open_boundary_transport = slim.DataStatic.create_constant(mesh, [0, 12])
open_boundaries = ["Open", "Open1"]

for boundary_name in open_boundaries:
    tracer_problem.set_open_boundary(
        boundary_name,
        tracer=open_boundary_tracer,
        eta=open_boundary_eta,
        u=open_boundary_transport,
        uv_is_transport=True,
        diffusion=False,
    )
    shallow_water.set_open_boundary(
        boundary_name,
        eta=open_boundary_eta,
        u=open_boundary_transport,
        uv_is_transport=True,
    )

tracer_problem.set_wall_boundary("Wall")
shallow_water.set_wall_boundary("Wall")

We define a helper function to export the solution to XDMF. In addition to the solution fields, we monitor the maximum change in tracer concentration since the previous export.

output_name = "advection"
mesh.create_xdmf(output_name)


def export_solution(export_index, t, state, previous_tracer_state):
    # transport = state.get_sub(0).get()
    # tracer_components = state.get_sub(1).get()

    # mesh.append_to_xdmf(
    #     output_name,
    #     export_index,
    #     t,
    #     [
    #         ("H", transport[0]),
    #         ("Hu", transport[1:]),
    #         ("C1", tracer_components[0]),
    #         ("C2", tracer_components[1]),
    #     ],
    # )

    # max_difference = slim4.mpi.max_scalar(
    #     np.abs(previous_tracer_state - tracer_components).max()
    # )
    # if slim4.mpi.rank == 0:
    #     print(f"Export {export_index}, DIFF : {max_difference:.3e}")

    # previous_tracer_state[:] = tracer_components
    return export_index + 1

Finally, we choose the time integration method and run the simulation. This example defaults to implicit Euler, but an explicit Runge-Kutta option is also kept for comparison.

implicit = False
dt = shallow_water.cfl
nsub = 100
t_end = 50000

if implicit:
    linear_system = tracer_problem.new_linear_system("")
    dt *= nsub

export_index = 0
t = 0
# previous_tracer_state = solution.get_sub(1).get().copy()
# export_index = export_solution(export_index, t, solution, previous_tracer_state)


from time import perf_counter as clock

ts = np.zeros(6)
iteration = 0
slim.synchronize()
ts[0] = clock()
while t < t_end:
    if implicit:
        t = slim.iterate_implicit_euler(tracer_problem, linear_system, solution, t, dt)
    else:
        t = slim.iterate_low_storage_runge_kutta(
            tracer_problem, solution, t, dt, 2/13, nsub
        )
    iteration += 1
    slim.synchronize()
    ts[iteration] = clock()
    # export_index = export_solution(export_index, t, solution, previous_tracer_state)
    print(iteration)
    if iteration == 5:
        break

dts = np.diff(ts)
print(dts)
print(np.mean(dts), np.std(dts), np.min(dts), np.max(dts))