Demo: Using the Block Solver for the Navier-Stokes Equations

This example demonstrates how to use the block solver capabilities in FLATiron to solve the incompressible Navier-Stokes equations. Here, we will split the velocity and pressure fields into separate blocks and solve them using a block preconditioner. Since this demo is intended to illustrate the block solver, we will point the reader to Demo: Transient Navier-Stokes for documentation on the set-up of the NSE problem.

Setting Up the Problem

Set up the Navier-Stokes problem as in Demo: Transient Navier-Stokes. Similar to the transient Navier-Stokes demo, the mesh is read from a located in demo/mesh/. The mesh can be generated using Gmsh and the provided bfs.geo file.

import dolfinx
import matplotlib.pyplot as plt
import numpy as np

from flatiron_tk.mesh import Mesh
from flatiron_tk.physics import TransientNavierStokes
from flatiron_tk.solver import BlockNonLinearSolver
from flatiron_tk.solver import BlockSplitTree
from flatiron_tk.solver import ConvergenceMonitor
from flatiron_tk.solver import NonLinearProblem
from mpi4py import MPI
from petsc4py import PETSc

# Define boundary condition functions
def no_slip(x):
    return np.stack((np.zeros(x.shape[1]), np.zeros(x.shape[1])))

def zero_pressure(x):
    return np.zeros(x.shape[1], dtype=dolfinx.default_scalar_type)

def inlet_velocity(x, t):
    values = np.zeros((2, x.shape[1]), dtype=dolfinx.default_scalar_type)
    values[0] = np.sin(t)
    return values

# Create the Mesh object by reading the mesh from file
mesh_file = '../mesh/bfs.msh'
mesh = Mesh(mesh_file=mesh_file)

# Create the TransientNavierStokes object
nse = TransientNavierStokes(mesh)
nse.set_element('CG', 1, 'CG', 1)
nse.build_function_space()

# Set physical parameters
dt = 0.05
mu = 0.001
rho = 1
nse.set_time_step_size(dt)
nse.set_midpoint_theta(0.5)
nse.set_density(rho)
nse.set_dynamic_viscosity(mu)

# Add stabilization
nse.set_weak_form(stab=True)

# Get function spaces to build boundary conditions
V_u = nse.get_function_space('u').collapse()[0]
V_p = nse.get_function_space('p').collapse()[0]

# Create functions for boundary conditions
inlet_v = dolfinx.fem.Function(V_u)

# At each time step:
t = 0.0
inlet_v.interpolate(lambda x: inlet_velocity(x, t))

zero_p = dolfinx.fem.Function(V_p)
zero_p.interpolate(zero_pressure)

zero_v = dolfinx.fem.Function(V_u)
zero_v.interpolate(no_slip)

# Define boundary conditions dictionary
u_bcs = {
        7: {'type': 'dirichlet', 'value': inlet_v},
        8: {'type': 'dirichlet', 'value': zero_v},
        10: {'type': 'dirichlet', 'value': zero_v},
        }

p_bcs = {
        9: {'type': 'dirichlet', 'value': zero_p},
        }

bc_dict = {'u': u_bcs,
        'p': p_bcs}

nse.set_bcs(bc_dict)

# Set the output writer
nse.set_writer('output', 'xdmf')

# Define the problem
problem = NonLinearProblem(nse)

We will now set up the block solver. We first need to define the block structure of our problem. In this case, we have two blocks: one for the velocity field and one for the pressure field. We set the preconditioner type and options for each block using a function called set_ksp_<split_name>.

# U block preconditioner parameters
def set_ksp_u(ksp):
    ksp.setType(PETSc.KSP.Type.FGMRES)
    ksp.setMonitor(ConvergenceMonitor("|----KSPU", verbose=True))
    ksp.setTolerances(max_it=3)
    ksp.pc.setType(PETSc.PC.Type.JACOBI)
    ksp.setUp()

# P block preconditioner parameters
def set_ksp_p(ksp):
    ksp.setType(PETSc.KSP.Type.FGMRES)
    ksp.setMonitor(ConvergenceMonitor("|--------KSPP", verbose=True))
    ksp.setTolerances(max_it=5)
    ksp.pc.setType(PETSc.PC.Type.HYPRE)
    ksp.pc.setHYPREType("boomeramg")
    ksp.setUp()

# Outer solver parameters
def set_outer_ksp(ksp):
    ksp.setType(PETSc.KSP.Type.FGMRES)
    ksp.setGMRESRestart(30)
    ksp.setTolerances(rtol=1e-100, atol=1e-10)
    ksp.setMonitor(ConvergenceMonitor("Outer ksp"))

We then define the block structure using BlockSplitTree and a dictionary that maps each block to its corresponding function in the Navier-Stokes problem.

# Define the block structure
split = {
        'fields': ('u', 'p'),
        'composite_type': 'schur',
        'schur_fact_type': 'full',
        'schur_pre_type': 'a11',
        'ksp0_set_function': set_ksp_u,
        'ksp1_set_function': set_ksp_p
        }

# Create the Block Solver
tree = BlockSplitTree(nse, splits=split)

Finally, we create the BlockNonLinearSolver object, passing in the tree structure, the communicator, and the outer solver parameters.

# Create the Block Non-Linear Solver
solver = BlockNonLinearSolver(tree, MPI.COMM_WORLD, problem, outer_ksp_set_function=set_outer_ksp)

Now we can solver the Navier-Stokes equations using the block solver as normal.

while t < 0.50:
    print(f'Solving at time t = {t:.2f}')

    # Set the inlet velocity for the current time step
    inlet_v.interpolate(lambda x: inlet_velocity(x, t))


    # Solve the problem
    solver.solve()

    nse.update_previous_solution()
    nse.write(time_stamp=t)

    # Update time
    t += dt

Full Script

import dolfinx
import matplotlib.pyplot as plt
import numpy as np

from flatiron_tk.mesh import Mesh
from flatiron_tk.physics import TransientNavierStokes
from flatiron_tk.solver import BlockNonLinearSolver
from flatiron_tk.solver import BlockSplitTree
from flatiron_tk.solver import ConvergenceMonitor
from flatiron_tk.solver import NonLinearProblem
from mpi4py import MPI
from petsc4py import PETSc

# Define boundary condition functions
def no_slip(x):
    return np.stack((np.zeros(x.shape[1]), np.zeros(x.shape[1])))

def zero_pressure(x):
    return np.zeros(x.shape[1], dtype=dolfinx.default_scalar_type)

def inlet_velocity(x, t):
    values = np.zeros((2, x.shape[1]), dtype=dolfinx.default_scalar_type)
    values[0] = np.sin(t)
    return values

# Create the Mesh object by reading the mesh from file
mesh_file = '../mesh/bfs.msh'
mesh = Mesh(mesh_file=mesh_file)

# Create the TransientNavierStokes object
nse = TransientNavierStokes(mesh)
nse.set_element('CG', 1, 'CG', 1)
nse.build_function_space()

# Set physical parameters
dt = 0.05
mu = 0.001
rho = 1
nse.set_time_step_size(dt)
nse.set_midpoint_theta(0.5)
nse.set_density(rho)
nse.set_dynamic_viscosity(mu)

# Add stabilization
nse.set_weak_form(stab=True)

# Get function spaces to build boundary conditions
V_u = nse.get_function_space('u').collapse()[0]
V_p = nse.get_function_space('p').collapse()[0]

# Create functions for boundary conditions
inlet_v = dolfinx.fem.Function(V_u)

# At each time step:
t = 0.0
inlet_v.interpolate(lambda x: inlet_velocity(x, t))

zero_p = dolfinx.fem.Function(V_p)
zero_p.interpolate(zero_pressure)

zero_v = dolfinx.fem.Function(V_u)
zero_v.interpolate(no_slip)

# Define boundary conditions dictionary
u_bcs = {
        7: {'type': 'dirichlet', 'value': inlet_v},
        8: {'type': 'dirichlet', 'value': zero_v},
        10: {'type': 'dirichlet', 'value': zero_v},
        }

p_bcs = {
        9: {'type': 'dirichlet', 'value': zero_p},
        }

bc_dict = {'u': u_bcs,
        'p': p_bcs}

nse.set_bcs(bc_dict)

# Set the output writer
nse.set_writer('output', 'xdmf')

# Define the problem
problem = NonLinearProblem(nse)

# Set up Block Solver
# U block preconditioner parameters
def set_ksp_u(ksp):
    ksp.setType(PETSc.KSP.Type.FGMRES)
    ksp.setMonitor(ConvergenceMonitor("|----KSPU", verbose=True))
    ksp.setTolerances(max_it=3)
    ksp.pc.setType(PETSc.PC.Type.JACOBI)
    ksp.setUp()

# P block preconditioner parameters
def set_ksp_p(ksp):
    ksp.setType(PETSc.KSP.Type.FGMRES)
    ksp.setMonitor(ConvergenceMonitor("|--------KSPP", verbose=True))
    ksp.setTolerances(max_it=5)
    ksp.pc.setType(PETSc.PC.Type.HYPRE)
    ksp.pc.setHYPREType("boomeramg")
    ksp.setUp()

# Outer solver parameters
def set_outer_ksp(ksp):
    ksp.setType(PETSc.KSP.Type.FGMRES)
    ksp.setGMRESRestart(30)
    ksp.setTolerances(rtol=1e-100, atol=1e-10)
    ksp.setMonitor(ConvergenceMonitor("Outer ksp"))

# Define the block structure
split = {
        'fields': ('u', 'p'),
        'composite_type': 'schur',
        'schur_fact_type': 'full',
        'schur_pre_type': 'a11',
        'ksp0_set_function': set_ksp_u,
        'ksp1_set_function': set_ksp_p
        }

# Create the Block Tree structure
tree = BlockSplitTree(nse, splits=split)

# Create the Block Nonlinear Solver
solver = BlockNonLinearSolver(tree, MPI.COMM_WORLD, problem, outer_ksp_set_function=set_outer_ksp)

while t < 10.0:
    print(f'Solving at time t = {t:.2f}')

    # Set the inlet velocity for the current time step
    inlet_v.interpolate(lambda x: inlet_velocity(x, t))


    # Solve the problem
    solver.solve()

    nse.update_previous_solution()
    nse.write(time_stamp=t)

    # Update time
    t += dt