Implementing method of lines using PETSc TS

[domi]

Hi,

I have a working example that solves the simple linear PDE \partial_t T - \nabla^2T - r = 0 using the method of lines, where petsc ts is used (see my mwe below):

import numpy as np
from mpi4py import MPI
from dolfinx import mesh, log
from dolfinx.io import XDMFFile
from dolfinx.fem import (functionspace, Function, locate_dofs_topological, dirichletbc, form)
import dolfinx.fem.petsc as petsc
from basix.ufl import element
from ufl import TestFunction, dx, grad, inner, TrialFunction
from petsc4py import PETSc

# 3d geometry parameters
box_length = 1
cells_number_per_dim = 4

# time stepping
t = 0
t_final = 10
dt_init = 1.

msh = mesh.create_box(MPI.COMM_WORLD,[[-box_length/2, -box_length/2, -box_length/2], [box_length/2,box_length/2, box_length/2]],
                      [cells_number_per_dim, cells_number_per_dim, cells_number_per_dim])

CG1_elem = element("Lagrange", msh.basix_cell(), 1)

V = functionspace(msh, CG1_elem)

T_test = TestFunction(V)
T = TrialFunction(V)

# Dirichlet boundary conditions on cube surface
tdim = msh.topology.dim
fdim = tdim - 1
msh.topology.create_connectivity(fdim, tdim)

boundary_conditions = [
    (lambda x: np.isclose(x[0], -box_length/2), 'dof_bdry_1'),
    (lambda x: np.isclose(x[0],  box_length/2), 'dof_bdry_2'),
    (lambda x: np.isclose(x[1], -box_length/2), 'dof_bdry_3'),
    (lambda x: np.isclose(x[1],  box_length/2), 'dof_bdry_4'),
    (lambda x: np.isclose(x[2], -box_length/2), 'dof_bdry_5'),
    (lambda x: np.isclose(x[2],  box_length/2), 'dof_bdry_6')
]

dof_boundaries = {}

for marker, dof_name in boundary_conditions:
    facets_bdry = mesh.locate_entities_boundary(msh, dim=fdim, marker=marker)
    dof_boundaries[dof_name] = locate_dofs_topological(V=V, entity_dim=fdim, entities=facets_bdry)


class T_exact_fun():
    def __init__(self, t):
        self.t = t

    def __call__(self, x):
        values = np.zeros((1, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = np.exp(-self.t) * (x[0]**2 + 2*x[1]**2)
        return values

class r_fun():
    def __init__(self, t):
        self.t = t

    def __call__(self, x):
        values = np.zeros((1, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = -np.exp(-self.t) * (x[0]**2 + 2*x[1]**2 + 6)
        return values

T_exact = Function(V)
T_exact.interpolate(T_exact_fun(t))

r = Function(V)
r.interpolate(r_fun(t))

bc1_fun = Function(V)
bc1_fun.interpolate(T_exact_fun(t))
bc1 = dirichletbc(bc1_fun, dof_boundaries['dof_bdry_1'])

bc2_fun = Function(V)
bc2_fun.interpolate(T_exact_fun(t))
bc2 = dirichletbc(bc2_fun, dof_boundaries['dof_bdry_2'])

bc3_fun = Function(V)
bc3_fun.interpolate(T_exact_fun(t))
bc3 = dirichletbc(bc3_fun, dof_boundaries['dof_bdry_3'])

bc4_fun = Function(V)
bc4_fun.interpolate(T_exact_fun(t))
bc4 = dirichletbc(bc4_fun, dof_boundaries['dof_bdry_4'])

bc5_fun = Function(V)
bc5_fun.interpolate(T_exact_fun(t))
bc5 = dirichletbc(bc5_fun, dof_boundaries['dof_bdry_5'])

bc6_fun = Function(V)
bc6_fun.interpolate(T_exact_fun(t))
bc6 = dirichletbc(bc6_fun, dof_boundaries['dof_bdry_6'])

Dirichlet_bcs = [bc1, bc2, bc3, bc4, bc5, bc6]

# Mass matrix
M_ufl = T_test * T * dx
M = petsc.assemble_matrix(form(M_ufl), Dirichlet_bcs)
M.assemble()

diffusion_ufl = inner(grad(T_test), grad(T)) * dx
diffusion = petsc.assemble_matrix(form(diffusion_ufl), Dirichlet_bcs)
diffusion.assemble()

reaction_ufl = - T_test * r * dx
reaction = petsc.assemble_vector(form(reaction_ufl))
reaction.assemble()

def calc_function(ts, t, T, T_dot, F):
    F.set(0)
    F += M * T_dot + diffusion * T - reaction


ts = PETSc.TS().create(comm=MPI.COMM_WORLD)
ts.setType('beuler')
ts.setIFunction(calc_function)

ts.setTime(t)
ts.setTimeStep(dt_init)
ts.setMaxTime(t_final)
ts.setExactFinalTime(PETSc.TS.ExactFinalTime.MATCHSTEP)
opts = PETSc.Options()
ts.setFromOptions()
T_vec = T_exact.x.petsc_vec.copy()

ts.solve(T_vec)

Now, I wonder how it is possible to extend this example, such that I can solve the nonlinear PDE \partial_t T - \nabla \cdot (T^2 +1) \nabla T - \hat{r} = 0.

Especially I am interested how I need to modify the diffusion_ufl expression, since TrialFunctions can only be used for linear PDEs. Also it seems that I would rather need a vector than a matrix in the nonlinear case. To be consistent, I also need to find a corresponding expression when defining the diffusion part in

def calc_function(ts, t, T, T_dot, F):
    F.set(0)
    F += M * T_dot + diffusion * T - reaction

Thanks in advance for your help.

[nate]

I wrote this a very long time ago with a very old version of DOLFINx, but maybe it can help you navigate the interface:

petsc4py-ts/dolfin_dg/dolfinx/ts.py

there are some demos here:

dolfin_dg/demo/dolfinx/ts at nate/petsc4py-ts · nate-sime/dolfin_dg

I have no idea how to get this up and running again, and I didn’t merge it with the main branch because I didn’t get the time to verify my implementation.

[domi]

Dear nate,

thank you, the code that you provided really helped me. However, as you already expected, I was not able to run it using dolfinx 0.9. In particular, the SystemAssembler class is not available anymore, i.e.

self.f_assembler = SystemAssembler(f_j, f, self.bcs)

does not work. I tried to look it up in the internet but was not able to find out what it exactly does and how it should be invoked in my case. I attached some code for my problem, where I tried to keep as close to the structure you provided:

import abc
import numpy as np
from mpi4py import MPI
from dolfinx import mesh
from dolfinx.fem import functionspace, Constant, Function, locate_dofs_topological, dirichletbc, assemble_vector
from basix.ufl import element
from ufl import derivative, dx, grad, inner, TestFunction
from petsc4py import PETSc


class PETScTSProblem(abc.ABC):

    def __init__(self):
        pass

    @abc.abstractmethod
    def F(self, u, u_t):
        pass

    def dFdu(self, u, u_t):
        return derivative(self.F(u, u_t), u)

    def dFdut(self, u, u_t):
        return derivative(self.F(u, u_t), u_t)


class PETScTSSolver:

    def __init__(self, ts_problem, u, bcs, tdexprs=[], mass_solver_prefix="mass_"):
        self.u = u
        self.V = u.ufl_function_space()
        self.comm = self.V.mesh.comm
        self.tdexprs = tdexprs
        self.mass_solver_prefix = mass_solver_prefix

        self.u_t = Function(self.V)
        self.sigma_t = Constant(self.V.mesh, 0.0)

        f = ts_problem.F(self.u, self.u_t)
        f_j = self.sigma_t * ts_problem.dFdut(self.u, self.u_t) + ts_problem.dFdu(self.u, self.u_t)

        self.bcs = bcs
        self.f_assembler = SystemAssembler(f_j, f, self.bcs)

        self.ts = PETSc.TS().create(self.comm)


    def evalIFunction(self, ts, t, x, xdot, F):
        for tdexpr in self.tdexprs:
            tdexpr.t = t
        self.u.x.set = x.array
        self.u_t.x.set = xdot.array
        self.f_assembler.assemble(F, self.u.x)

    def evalIJacobian(self, ts, t, x, xdot, sigma_t, A, P):
        for tdexpr in self.tdexprs:
            tdexpr.t = t
        self.u.x.set = x.array
        self.u_t.x.set = xdot.array
        self.sigma_t.assign(sigma_t)
        self.f_assembler.assemble(A, self.u.vector)

    def set_from_options(self):
        self.ts.setFromOptions()
        self.ts.getSNES().setFromOptions()

    def solve(self, t0=0.0, dt=0.1, max_t=1.0, max_steps=10):
        # Initialise tensors with correct sparsity pattern
        #self.f_assembler.assemble(self.FJ, self.F)

        # Solution vector
        w = Function(self.V)
        w.x.array[:] = self.u.x.array[:]
        w_PETScVEC = PETSc.Vec().createWithArray(w.x.array[:])
        self.ts.setIFunction(self.evalIFunction)#, self.F.vec())
        self.ts.setIJacobian(self.evalIJacobian)#, J=self.FJ.mat(), P=self.FJ.mat())
        self.ts.setSolution(w_PETScVEC)
        self.ts.setTime(t0)
        self.ts.setTimeStep(dt)
        self.ts.setMaxTime(max_t)
        self.ts.setMaxSteps(max_steps)
        self.ts.solve(w_PETScVEC)
        self.u.x.array[:] = w.x.array[:]


# 3d geometry parameters
box_length = 1
cells_number_per_dim = 4

t = 0

msh = mesh.create_box(MPI.COMM_WORLD,[[-box_length/2, -box_length/2, -box_length/2], [box_length/2,box_length/2, box_length/2]],
                      [cells_number_per_dim, cells_number_per_dim, cells_number_per_dim])

CG1_elem = element("Lagrange", msh.basix_cell(), 1)

V = functionspace(msh, CG1_elem)

u_test = TestFunction(V)
u = Function(V)

# Dirichlet boundary conditions on cube surface
tdim = msh.topology.dim
fdim = tdim - 1
msh.topology.create_connectivity(fdim, tdim)

boundary_conditions = [
    (lambda x: np.isclose(x[0], -box_length/2), 'dof_bdry_1')
]

dof_boundaries = {}

for marker, dof_name in boundary_conditions:
    facets_bdry = mesh.locate_entities_boundary(msh, dim=fdim, marker=marker)
    dof_boundaries[dof_name] = locate_dofs_topological(V=V, entity_dim=fdim, entities=facets_bdry)


class T_exact_fun():
    def __init__(self, t):
        self.t = t

    def __call__(self, x):
        values = np.zeros((1, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = np.exp(-self.t) * (x[0]**2 + 2*x[1]**2)
        return values


class r_fun():
    def __init__(self, t):
        self.t = t

    def __call__(self, x):
        values = np.zeros((1, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = -np.exp(-self.t) * (x[0]**2 + 2*x[1]**2 + 6)
        return values

class HeatProblem(PETScTSProblem):
    def F(self, u, u_t):
        return u_t * u_test * dx + inner(grad(u), grad(u_test)) * dx - r * u_test * dx

T_exact = Function(V)
T_exact.interpolate(T_exact_fun(t))

r = Function(V)
r.interpolate(r_fun(t))

bc1_fun = Function(V)
bc1_fun.interpolate(T_exact_fun(t))
bc1 = dirichletbc(bc1_fun, dof_boundaries['dof_bdry_1'])

Dirichlet_bcs = [bc1]

ts = PETScTSSolver(HeatProblem(), u, Dirichlet_bcs)
ts.solve(t0=0.0, dt=1e-2, max_t=1.0, max_steps=1001)

Could you give me some guidance what needs to be done in order to substitute SystemAssembler correctly using dolfinx 0.9?

Thank you in advance.

[domi]

I made some progress and I think that the remaining issue is that I am not able to properly enforce Dirichlet Boundary conditions on resiual and solution vector in evalIFunction method. Here is what I have so far:

import numpy as np
from mpi4py import MPI
from dolfinx import mesh, log
from dolfinx.io import XDMFFile
from dolfinx.fem import (functionspace, Constant, Function, locate_dofs_topological, dirichletbc, DirichletBC,
                         Expression, assemble_scalar, assemble_vector, apply_lifting, set_bc, form)
import dolfinx.fem.petsc as petsc
from basix.ufl import element
from ufl import derivative, dx, div, grad, inner, SpatialCoordinate, TrialFunction, TestFunction
from petsc4py import PETSc

dt = 0.01

class PETScTSSolver:
    # F(t, u, u') = G(t, u)
    def __init__(self, ts_problem, u, bcs):
        self.u = u
        self.V = u.ufl_function_space()
        self.comm = self.V.mesh.comm
        self.u_t = Function(self.V)
        self.sigma_t = Constant(self.V.mesh, 1/dt)

        self.f = ts_problem.F(self.u, self.u_t)
        self.f_j = self.sigma_t * ts_problem.dFdut(self.u, self.u_t) + ts_problem.dFdu(self.u, self.u_t)

        self.bcs = bcs
        self.ts = PETSc.TS().create(self.comm)

    def evalIFunction(self, ts, t, x, xdot, F):
        # PETScTS implicit callback
        x.copy(self.u.x.petsc_vec)
        xdot.copy(self.u_t.x.petsc_vec)
        F.set(0)
        petsc.assemble_vector(F, form(self.f))

        # Apply boundary condition
        apply_lifting(F, [form(self.f)], bcs=[Dirichlet_bcs], x0=[self.u.x.petsc_vec], alpha=-1.0)
        F.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        set_bc(F, [Dirichlet_bcs], self.u.x.petsc_vec, -1.0)

        F.assemble()

        return True


    def evalIJacobian(self, ts, t, x, xdot, sigma_t, J, P):
        #x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        x.copy(self.u.x.petsc_vec)
        #self.u.x.petsc_vec.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
        xdot.copy(self.u_t.x.petsc_vec)
        J.zeroEntries()
        petsc.assemble_matrix(J, form(self.f_j))
        J.assemble()
        return True

    def set_from_options(self):
        self.ts.setFromOptions()
        self.ts.getSNES().setFromOptions()

    def solve(self, t0=0.0, dt=0.1, max_t=1.0, max_steps=10):

        # Solution vector
        w = Function(self.V)
        w.x.array[:] = T_exact.x.petsc_vec #self.u.x.petsc_vec

        self.ts.setIFunction(self.evalIFunction)#, self.F.vec())
        self.ts.setIJacobian(self.evalIJacobian)#, J=self.FJ.mat(), P=self.FJ.mat())
        self.ts.setSolution(w.x.petsc_vec)
        self.ts.setTime(t0)
        self.ts.setTimeStep(dt)
        self.ts.setMaxTime(max_t)
        self.ts.setMaxSteps(max_steps)
        self.ts.solve(w.x.petsc_vec)
        self.u.x.array[:] = w.x.petsc_vec


# 3d geometry parameters
box_length = 1
cells_number_per_dim = 4

t = 0

# diagnostics
output_folder = './output'

msh = mesh.create_box(MPI.COMM_WORLD,[[-box_length/2, -box_length/2, -box_length/2], [box_length/2,box_length/2, box_length/2]],
                      [cells_number_per_dim, cells_number_per_dim, cells_number_per_dim])

CG1_elem = element("Lagrange", msh.basix_cell(), 1)

V = functionspace(msh, CG1_elem)

u_test = TestFunction(V)
u = Function(V)

# Dirichlet boundary conditions on cube surface
tdim = msh.topology.dim
fdim = tdim - 1
msh.topology.create_connectivity(fdim, tdim)

with XDMFFile(msh.comm, output_folder + "/T_sol_MoL.xdmf", "w") as T_sol_plot:
    T_sol_plot.write_mesh(msh)
with XDMFFile(msh.comm, output_folder + "/T_exact_MoL.xdmf", "w") as T_exact_plot:
    T_exact_plot.write_mesh(msh)

num_dofs = V.dofmap.index_map.size_local
total_dofs = V.dofmap.index_map.size_global

boundary_conditions = [
    (lambda x: np.isclose(x[0], -box_length/2), 'dof_bdry_1'),
    (lambda x: np.isclose(x[0],  box_length/2), 'dof_bdry_2'),
    (lambda x: np.isclose(x[1], -box_length/2), 'dof_bdry_3'),
    (lambda x: np.isclose(x[1],  box_length/2), 'dof_bdry_4'),
    (lambda x: np.isclose(x[2], -box_length/2), 'dof_bdry_5'),
    (lambda x: np.isclose(x[2],  box_length/2), 'dof_bdry_6')
]

dof_boundaries = {}

for marker, dof_name in boundary_conditions:
    facets_bdry = mesh.locate_entities_boundary(msh, dim=fdim, marker=marker)
    dof_boundaries[dof_name] = locate_dofs_topological(V=V, entity_dim=fdim, entities=facets_bdry)


class T_exact_fun():
    def __init__(self, t):
        self.t = t

    def __call__(self, x):
        values = np.zeros((1, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = np.exp(-self.t) * (x[0]**2 + 2*x[1]**2)
        return values


class r_fun():
    def __init__(self, t):
        self.t = t

    def __call__(self, x):
        values = np.zeros((1, x.shape[1]), dtype=PETSc.ScalarType)
        values[0] = -np.exp(-self.t) * (x[0]**2 + 2*x[1]**2 + 6)
        return values

def T_0(x):
    values = np.zeros((1, x.shape[1]))
    values[0] = 0
    return values

class HeatProblem():
    def F(self, u, u_t):
        return u_t * u_test * dx + inner(grad(u), grad(u_test)) * dx - r * u_test * dx

    def dFdu(self, u, u_t):
        return derivative(self.F(u, u_t), u)

    def dFdut(self, u, u_t):
        return derivative(self.F(u, u_t), u_t)

    def G(self, u):
        pass

    def Gu(self, u):
        return derivative(self.G(u), u)



T_exact = Function(V)
T_exact.interpolate(T_exact_fun(t))

r = Function(V)
r.interpolate(r_fun(t))

bc1_fun = Function(V)
bc1_fun.interpolate(T_exact_fun(t))
bc1 = dirichletbc(bc1_fun, dof_boundaries['dof_bdry_1'])

bc2_fun = Function(V)
bc2_fun.interpolate(T_exact_fun(t))
bc2 = dirichletbc(bc2_fun, dof_boundaries['dof_bdry_2'])

bc3_fun = Function(V)
bc3_fun.interpolate(T_exact_fun(t))
bc3 = dirichletbc(bc3_fun, dof_boundaries['dof_bdry_3'])

bc4_fun = Function(V)
bc4_fun.interpolate(T_exact_fun(t))
bc4 = dirichletbc(bc4_fun, dof_boundaries['dof_bdry_4'])

bc5_fun = Function(V)
bc5_fun.interpolate(T_exact_fun(t))
bc5 = dirichletbc(bc5_fun, dof_boundaries['dof_bdry_5'])

bc6_fun = Function(V)
bc6_fun.interpolate(T_exact_fun(t))
bc6 = dirichletbc(bc6_fun, dof_boundaries['dof_bdry_6'])

Dirichlet_bcs = [bc1, bc2, bc3, bc4, bc5, bc6]

ts = PETScTSSolver(HeatProblem(), u, Dirichlet_bcs)
ts.solve(t0=0.0, dt=dt, max_t=1.0)

However, running this code causes segmentation faults due to

 apply_lifting(F, [form(self.f)], bcs=[Dirichlet_bcs], x0=[self.u.x.petsc_vec], alpha=-1.0)

[a-lyrae]

I believe you should use the Jacobian and the current solution x when you use apply_lifting. Here you are passing the residual and the initial guess u…