from itertools import chain

from pystencils.session import *
from lbmpy.session import *
from lbmpy.methods.default_moment_sets import cascaded_moment_sets_literature
from lbmpy.moments import moments_up_to_component_order, exponent_tuple_sort_key, MOMENT_SYMBOLS
from lbmpy.methods import DensityVelocityComputation, create_from_equilibrium, CollisionSpaceInfo
from lbmpy.equilibrium import GenericDiscreteEquilibrium, ContinuousHydrodynamicMaxwellian

from pystencils.sympyextensions import scalar_product as dot

Demo: Lattice Boltzmann Methods for the Shallow Water Equations

In this notebook we present how lbmpy’s low-level method definition interface may be used to rapidly set up two lattice Boltzmann models for solving the shallow water equations (SWE). The shallow water equations approximate the Navier-Stokes equations in a regime where the vertical length scale of the flow is much smaller than the horizontal length scales. Velocity is thus averaged across the third dimension, and density is replaced by water column height \(h\).

We reconstruct the central moment-based method presented by De Rosis in https://doi.org/10.1016/j.cma.2017.03.001 and the cumulant-based method of Venturi et al. (https://doi.org/10.1016/j.advwatres.2019.103474).

d2q9 = LBStencil(Stencil.D2Q9)

# gravitational acceleration
g = sp.Symbol('g')

# water column height
h = sp.Symbol('h')

# velocity
u = sp.symbols(f'u_:{d2q9.D}')

# relaxation rate
ω_s = sp.Symbol('omega_s')

# moment variables
x, y, _ =  MOMENT_SYMBOLS

Discrete Shallow Water Equilibrium

The central moment-based method makes use of the discrete shallow water LB equilibrium, which was originally defined by Zhou et al. in https://doi.org/10.1016/s0045-7825(02)00291-8. We implement it by listing equations for its components, and creating an instance of lbmpy.equilibrium.GenericDiscreteEquilibrium:

def f_eq(ξ):
    ξ_sum = sum(abs(ξ_i) for ξ_i in ξ)
    if ξ_sum == 0:
        return h * (1 - (5 * g * h) / 6 - (2 * dot(u,u)) / 3)
    else:
        λ = 1 if ξ_sum == 1 else sp.Rational(1, 4)
        common = (g * h) / 6 + dot(ξ, u) / 3 + dot(ξ, u)**2 / 2 - dot(u,u) / 6
        return λ * h * common

eq_populations = [f_eq(ξ) for ξ in d2q9]
discrete_swe_eq = GenericDiscreteEquilibrium(d2q9, eq_populations, h, u)

discrete_swe_eq

Instance of GenericDiscreteEquilibrium
Discrete Populations:
$f_0 = h \left(- \frac{5 g h}{6} - \frac{2 u_{0}^{2}}{3} - \frac{2 u_{1}^{2}}{3} + 1\right)$
$f_1 = h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} + \frac{u_{1}^{2}}{3} + \frac{u_{1}}{3}\right)$
$f_2 = h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} + \frac{u_{1}^{2}}{3} - \frac{u_{1}}{3}\right)$
$f_3 = h \left(\frac{g h}{6} + \frac{u_{0}^{2}}{3} - \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6}\right)$
$f_4 = h \left(\frac{g h}{6} + \frac{u_{0}^{2}}{3} + \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6}\right)$
$f_5 = \frac{h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} - \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6} + \frac{u_{1}}{3} + \frac{\left(- u_{0} + u_{1}\right)^{2}}{2}\right)}{4}$
$f_6 = \frac{h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} + \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6} + \frac{u_{1}}{3} + \frac{\left(u_{0} + u_{1}\right)^{2}}{2}\right)}{4}$
$f_7 = \frac{h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} - \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6} - \frac{u_{1}}{3} + \frac{\left(- u_{0} - u_{1}\right)^{2}}{2}\right)}{4}$
$f_8 = \frac{h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} + \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6} - \frac{u_{1}}{3} + \frac{\left(u_{0} - u_{1}\right)^{2}}{2}\right)}{4}$

Continuous Shallow Water Equilibrium

The cumulant-based method uses a special form of the continuous Maxwellian equilibrium, which is adapted to the shallow water equations by redefining its squared speed of sound parameter to be \(c_s^2 = gh / 2\):

c_s_sq = g * h / 2
cont_sqe_eq = ContinuousHydrodynamicMaxwellian(dim=d2q9.D, rho=h, c_s_sq=c_s_sq)

cont_sqe_eq

Continuous Hydrodynamic Maxwellian Equilibrium			$f (h, \left( u_{0}, \ u_{1}\right), \left( v_{0}, \ v_{1}\right)) = \frac{e^{\frac{- \left(- u_{0} + v_{0}\right)^{2} - \left(- u_{1} + v_{1}\right)^{2}}{g h}}}{\pi g}$
Compressible: ✓	Deviation Only: ✗	Order: ∞

Method Definition

The remaining parts of the method definition are nearly identical between both methods. Merely the collision spaces must be specified separately: CollisionSpace.CENTRAL_MOMENTS vs. CollisionSpace.CUMULANTS.

The on-line version of this notebook runs the following simulation using the central moment-based method. To activate the cumulant-based method instead, flip the switch in the next cell:

# toggle central moment / cumulant method

cumulant_method = False

polys = (sp.Integer(1), x, y, x * y, x**2 - y**2, x**2 + y**2, x * y**2, y * x**2, x**2 * y**2)
r_rates = [0, 0, 0, ω_s, ω_s, 1, 1, 1, 1]
rr_dict = dict(zip(polys, r_rates))

cqc = DensityVelocityComputation(d2q9, True, False,
                                 density_symbol=h,
                                 density_deviation_symbol=sp.Symbol('delta_h'))

if cumulant_method:
    swe_eq = cont_sqe_eq
    cspace = CollisionSpaceInfo(CollisionSpace.CUMULANTS)
else:
    swe_eq = discrete_swe_eq
    cspace = CollisionSpaceInfo(CollisionSpace.CENTRAL_MOMENTS)

swe_method = create_from_equilibrium(d2q9, swe_eq, cqc, rr_dict, cspace)

swe_method

Central-Moment-Based Method			Stencil: D2Q9	Zero-Centered Storage: ✗	Force Model: None

Instance of GenericDiscreteEquilibrium
Discrete Populations:
$f_0 = h \left(- \frac{5 g h}{6} - \frac{2 u_{0}^{2}}{3} - \frac{2 u_{1}^{2}}{3} + 1\right)$
$f_1 = h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} + \frac{u_{1}^{2}}{3} + \frac{u_{1}}{3}\right)$
$f_2 = h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} + \frac{u_{1}^{2}}{3} - \frac{u_{1}}{3}\right)$
$f_3 = h \left(\frac{g h}{6} + \frac{u_{0}^{2}}{3} - \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6}\right)$
$f_4 = h \left(\frac{g h}{6} + \frac{u_{0}^{2}}{3} + \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6}\right)$
$f_5 = \frac{h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} - \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6} + \frac{u_{1}}{3} + \frac{\left(- u_{0} + u_{1}\right)^{2}}{2}\right)}{4}$
$f_6 = \frac{h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} + \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6} + \frac{u_{1}}{3} + \frac{\left(u_{0} + u_{1}\right)^{2}}{2}\right)}{4}$
$f_7 = \frac{h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} - \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6} - \frac{u_{1}}{3} + \frac{\left(- u_{0} - u_{1}\right)^{2}}{2}\right)}{4}$
$f_8 = \frac{h \left(\frac{g h}{6} - \frac{u_{0}^{2}}{6} + \frac{u_{0}}{3} - \frac{u_{1}^{2}}{6} - \frac{u_{1}}{3} + \frac{\left(u_{0} - u_{1}\right)^{2}}{2}\right)}{4}$

Relaxation Info
Central Moment	Eq. Value	Relaxation Rate
$1$	$h$	$0$
$x$	$0$	$0$
$y$	$0$	$0$
$x y$	$0$	$\omega_{s}$
$x^{2} - y^{2}$	$0$	$\omega_{s}$
$x^{2} + y^{2}$	$g h^{2}$	$1$
$x y^{2}$	$- \frac{g h^{2} u_{0}}{2} - h u_{0} u_{1}^{2} + \frac{h u_{0}}{3}$	$1$
$x^{2} y$	$- \frac{g h^{2} u_{1}}{2} - h u_{0}^{2} u_{1} + \frac{h u_{1}}{3}$	$1$
$x^{2} y^{2}$	$\frac{g h^{2} u_{0}^{2}}{2} + \frac{g h^{2} u_{1}^{2}}{2} + \frac{g h^{2}}{6} + 3 h u_{0}^{2} u_{1}^{2} - \frac{h u_{0}^{2}}{3} - \frac{h u_{1}^{2}}{3}$	$1$

Circular Dam Break Scenario

We use the method derived above to simulate a circular dam break scenario. We discretize a square domain of \(40 \times 40\) meters using \(100 \times 100\) lattice cells, with periodic boundary conditions. A water column of \(2.5\) meters heights, and \(2.5\) meters radius, is placed in the middle. Otherwise, water depth is uniformly \(0.5\) meters. We neglect any friction due to the bed. We assume default Earth gravitational pull of \(9.81 \mathrm{m}/\mathrm{s}\). The simulated time step length is \(0.05\) seconds per time step.

L_lattice = 100 # LU
L_phys = 40 # meters
L_factor = L_phys / L_lattice # meters / lu

dt = 0.05 # seconds / step
g_phys = 9.81 # meters / second**2

g_lattice = g_phys / L_factor * dt**2 # lu / step**2
shear_rr = 1.1

dh = ps.create_data_handling((L_lattice, L_lattice), periodicity=True,
                             default_target=ps.Target.CPU)
h_field = dh.add_array('h', values_per_cell=1)
u_field = dh.add_array('u', values_per_cell=d2q9.D)
pdf_field = dh.add_array('f', values_per_cell=d2q9.Q)
pdf_tmp_field = dh.add_array('f_tmp', values_per_cell=d2q9.Q)

from lbmpy.macroscopic_value_kernels import (
    macroscopic_values_setter, macroscopic_values_getter)

setter = macroscopic_values_setter(swe_method,
                                   h_field.center, u_field.center_vector,
                                   pdf_field,
                                   set_pre_collision_pdfs=True)

setter_kernel = ps.create_kernel(setter).compile()

getter = macroscopic_values_getter(swe_method,
                                   h_field.center, u_field.center_vector,
                                   pdf_field,
                                   use_pre_collision_pdfs=True)
getter_kernel = ps.create_kernel(getter).compile()

periodic_sync = dh.synchronization_function(pdf_field.name)

lbm_config = LBMConfig(lb_method=swe_method)
opt = LBMOptimisation(pre_simplification=True, simplification=True,
                      symbolic_field=pdf_field, symbolic_temporary_field=pdf_tmp_field)
lb_kernel = create_lb_function(lbm_config=lbm_config, lbm_optimisation=opt)

Initialization of PDF field

def init():
    dh.fill(u_field.name, 0.0)

    h_arr = dh.cpu_arrays[h.name][1:-1,1:-1]

    x_c = y_c = L_phys / 2
    R = 2.5
    H_phys = 2.5
    b_phys = 0.5

    # set up column
    for i in range(100):
        for j in range(100):
            x = L_factor * (i + 0.5)
            y = L_factor * (j + 0.5)

            if (x - x_c)**2 + (y - y_c)**2 <= R**2:
                h_arr[i,j] = H_phys / L_factor
            else:
                h_arr[i,j] = b_phys / L_factor

    dh.run_kernel(setter_kernel, g=g_lattice)

Time Step

def step(output=False):
    periodic_sync()
    if output:
        dh.run_kernel(getter_kernel)
    dh.run_kernel(lb_kernel, g=g_lattice, omega_s=shear_rr)
    dh.swap(pdf_field.name, pdf_tmp_field.name)

    h_arr = dh.cpu_arrays[h.name][1:-1,1:-1]
    return L_factor * h_arr

Simulation

frames = 200
timesteps = 100
steps_per_frame = (timesteps//frames) + 1

def run():
    for _ in range(steps_per_frame - 1):
        step()
    return step(output=True)

init()

animation = plt.surface_plot_animation(run, frames=frames, zlim=(0, 3))
set_display_mode('video')
res = display_animation(animation)
res