Issue with `multiphenicsx.fem.petsc.set_bc` x0 argument

[ordinary-slim]

I can’t seem to figure out how to correctly set the x0 argument of multiphenicsx.fem.set_bc and multiphenicsx.fem.apply_lifting.

In this MWE I solve incrementally a linear heat equation on the left side of the unit square with Dirichlet heating on the “outside” boundary and the middle boundary is insulated:

from dolfinx import fem, mesh, io
import ufl
import numpy as np
from mpi4py import MPI
import multiphenicsx
import multiphenicsx.fem.petsc
import petsc4py
from petsc4py import PETSc

comm = MPI.COMM_WORLD
rank = comm.Get_rank()

def main():
    # Mesh and problems
    points_per_side = 8
    domain  = mesh.create_unit_square(MPI.COMM_WORLD, points_per_side, points_per_side, mesh.CellType.quadrilateral)
    tdim = domain.topology.dim
    CG1  = fem.functionspace(domain, ("Lagrange", 1),)
    DG0  = fem.functionspace(domain, ("Discontinuous Lagrange", 0),)
    domain.topology.create_entities(tdim)
    domain.topology.create_entities(tdim-1)
    domain.topology.create_connectivity(tdim,tdim)
    cell_map = domain.topology.index_map(tdim)

    # Activation
    active_els = fem.locate_dofs_geometrical(DG0, lambda x : x[0] <= 0.5 )
    active_els_func = fem.Function(DG0,name="active_els")
    active_els_func.x.array[active_els] = 1.0
    active_dofs = fem.locate_dofs_topological(CG1, tdim, active_els,remote=True)
    restriction = multiphenicsx.fem.DofMapRestriction(CG1.dofmap, active_dofs)

    # Parameters
    dt = fem.Constant(domain, 0.1)
    k   = fem.Constant(domain, 1.0)
    c_p = fem.Constant(domain, 1.0)
    rho = fem.Constant(domain, 1.0)
    T_hot = 2.0
    T_cold = 0.0
    num_timesteps = 5

    unp1 = fem.Function(CG1, name="unp1")
    un   = fem.Function(CG1, name="un")
    udc  = fem.Function(CG1, name="dirichlet")
    du   = fem.Function(CG1, name="delta_uh")
    alg_rhs = fem.Function(CG1, name="rhs")

    # Dirichlet BC
    def left_marker_dirichlet(x):
        return np.logical_or( np.isclose(x[1],1), np.logical_or(
                np.isclose(x[0],0), np.isclose(x[1],0)) )
    dirichlet_dofs = fem.locate_dofs_geometrical(CG1,left_marker_dirichlet)
    dirichlet_bcs = [fem.dirichletbc(udc, dirichlet_dofs)]
    
    # Forms
    local_active_els = active_els_func.x.array.nonzero()[0][:np.searchsorted(active_els_func.x.array.nonzero()[0], cell_map.size_local)]
    subdomain_idx = 1
    dx = ufl.Measure("dx", subdomain_data=[(subdomain_idx,local_active_els)])
    (u, v) = (ufl.TrialFunction(CG1), ufl.TestFunction(CG1))
    a_ufl = k*ufl.dot(ufl.grad(u), ufl.grad(v))*dx(subdomain_idx)
    a_ufl += (rho*c_p/dt)*u*v*dx(subdomain_idx)
    l_ufl = (rho*c_p/dt)*un*v*dx(subdomain_idx)

    a_compiled = fem.form(a_ufl)
    l_compiled = fem.form(l_ufl)

    '''
    # Normal solve
    A = multiphenicsx.fem.petsc.create_matrix(a_compiled, (restriction, restriction),)
    L = multiphenicsx.fem.petsc.create_vector(l_compiled, restriction)
    x = multiphenicsx.fem.petsc.create_vector(l_compiled, restriction=restriction)

    time = 0.0
    unp1.x.array[:] = T_cold
    udc.x.array[:] = T_hot
    vtk_writer = io.VTKFile(domain.comm, f"post/normal.pvd", "wb")
    for titer in range(num_timesteps):
        time += dt.value
        un.x.array[:] = unp1.x.array[:]
        # ASSEMBLE
        A.zeroEntries()
        multiphenicsx.fem.petsc.assemble_matrix(A,
                                                a_compiled,
                                                bcs=dirichlet_bcs,
                                                restriction=(restriction, restriction))
        A.assemble()
        with L.localForm() as l_local:
            l_local.set(0.0)
        multiphenicsx.fem.petsc.assemble_vector(L,
                                                l_compiled,
                                                restriction=restriction,)
        multiphenicsx.fem.petsc.apply_lifting(L, [a_compiled], [dirichlet_bcs], restriction=restriction,)
        L.ghostUpdate(addv=petsc4py.PETSc.InsertMode.ADD, mode=petsc4py.PETSc.ScatterMode.REVERSE)
        multiphenicsx.fem.petsc.set_bc(L,dirichlet_bcs,restriction=restriction)
        # LINEAR SOLVE
        opts = {"pc_type" : "lu", "pc_factor_mat_solver_type" : "mumps",}
        with x.localForm() as x_local:
            x_local.set(0.0)
        ksp = petsc4py.PETSc.KSP()
        ksp.create(domain.comm)
        ksp.setOperators(A)
        ksp_opts = PETSc.Options()
        for key,value in opts.items():
            ksp_opts[key] = value
        ksp.setFromOptions()
        ksp.solve(L, x)
        x.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
        ksp.destroy()
        # Unrestrict solution
        with unp1.x.petsc_vec.localForm() as usub_vector_local, \
                multiphenicsx.fem.petsc.VecSubVectorWrapper(x, CG1.dofmap, restriction) as x_wrapper:
                    usub_vector_local[:] = x_wrapper
        unp1.x.scatter_forward()

        vtk_writer.write_function([unp1,un,udc,active_els_func],t=time)
    vtk_writer.close()
    for la_ds in [A, L, x]:
        la_ds.destroy()
    '''

    # Incremental solve
    (u, v) = (unp1, ufl.TestFunction(CG1))
    a_ufl = k*ufl.dot(ufl.grad(u), ufl.grad(v))*dx(subdomain_idx)
    a_ufl += (rho*c_p/dt)*u*v*dx(subdomain_idx)
    l_ufl = (rho*c_p/dt)*un*v*dx(subdomain_idx)
    r_ufl = a_ufl - l_ufl

    j_ufl = ufl.derivative(a_ufl,unp1)
    j_compiled = fem.form(j_ufl)
    r_compiled = fem.form(r_ufl)
    A = multiphenicsx.fem.petsc.create_matrix(j_compiled, (restriction, restriction),)
    L = multiphenicsx.fem.petsc.create_vector(r_compiled, restriction)
    x = multiphenicsx.fem.petsc.create_vector(r_compiled, restriction=restriction)

    time = 0.0
    unp1.x.array[:] = T_cold
    udc.x.array[:] = T_hot
    vtk_writer = io.VTKFile(domain.comm, f"post/incremenetal.pvd", "wb")
    for titer in range(num_timesteps):
        time += dt.value
        un.x.array[:] = unp1.x.array[:]
        # ASSEMBLE
        A.zeroEntries()
        multiphenicsx.fem.petsc.assemble_matrix(A,
                                                j_compiled,
                                                bcs=dirichlet_bcs,
                                                restriction=(restriction, restriction))
        A.assemble()
        with L.localForm() as l_local:
            l_local.set(0.0)
        multiphenicsx.fem.petsc.assemble_vector(L,
                                                r_compiled,
                                                restriction=restriction,)
        L.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
        L.scale(-1)

        # Dirichlet
        multiphenicsx.fem.petsc.apply_lifting(L, [j_compiled], [dirichlet_bcs], x0=[unp1.x.petsc_vec], restriction=restriction)
        multiphenicsx.fem.petsc.set_bc(L,dirichlet_bcs, x0=unp1.x.petsc_vec, restriction=restriction)
        L.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)

        # RHS to function
        with alg_rhs.x.petsc_vec.localForm() as lfunsub_vector_local, \
                multiphenicsx.fem.petsc.VecSubVectorWrapper(L, CG1.dofmap, restriction) as l_wrapper:
                    lfunsub_vector_local[:] = l_wrapper
        alg_rhs.x.scatter_forward()


        # LINEAR SOLVE
        opts = {"pc_type" : "lu", "pc_factor_mat_solver_type" : "mumps",}
        with x.localForm() as x_local:
            x_local.set(0.0)
        ksp = petsc4py.PETSc.KSP()
        ksp.create(domain.comm)
        ksp.setOperators(A)
        ksp_opts = PETSc.Options()
        for k,v in opts.items():
            ksp_opts[k] = v
        ksp.setFromOptions()
        ksp.solve(L, x)
        x.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
        ksp.destroy()
        # Unrestrict solution
        with du.x.petsc_vec.localForm() as dusub_vector_local, \
                multiphenicsx.fem.petsc.VecSubVectorWrapper(x, CG1.dofmap, restriction) as x_wrapper:
                    dusub_vector_local[:] = x_wrapper
        du.x.scatter_forward()
        unp1.x.array[:] += du.x.array[:]

        vtk_writer.write_function([unp1,un,udc,active_els_func,alg_rhs],t=time)
    vtk_writer.close()
    for la_ds in [A, L, x]:
        la_ds.destroy()

if __name__=="__main__":
    main()

There is a a “normal” solve commented out and an incremental solve (residual, jacobian) both using multiphenicsx. First two time-steps of both:

At the first time-step, both solves obtain the same thing. At the second time-step, my incremental solve doesn’t apply correctly the Dirichlet BC. The RHS at the second time-step should be zero at the Dirichlet nodes after calling multiphenicsx.fem.petsc.set_bc but it is not in my case. It is mostly zero except at some nodes:

multiphenicsx.fem.petsc.set_bc(L,dirichlet_bcs, x0=unp1.x.petsc_vec, restriction=restriction)

I’m working with the latest commits of dolfinx, multiphenicsx, etc. on the test-env docker container. Any help is appreciated.

@francesco-ballarin please could you have a look?

[ordinary-slim]

Upon further testing it seems the mistake occurs already at the first time-step if the Dirichlet nodes start out already at their Dirichlet value in the solution u^{n+1}.

Moreover the problem does seem to be related to multiphenicsx since I saw no issues using the same approach with pure dolfinx. The following dolfinx example is the same as in the previous post except top, left and bottom boundaries are Dirichlet and right boundary is insulated.

from dolfinx import fem, mesh, io
import ufl
import numpy as np
from mpi4py import MPI
import petsc4py
from petsc4py import PETSc
import dolfinx.fem.petsc

comm = MPI.COMM_WORLD
rank = comm.Get_rank()

def main():
    # Mesh and problems
    points_per_side = 8
    domain  = mesh.create_unit_square(MPI.COMM_WORLD, points_per_side, points_per_side, mesh.CellType.quadrilateral)
    tdim = domain.topology.dim
    CG1  = fem.functionspace(domain, ("Lagrange", 1),)
    DG0  = fem.functionspace(domain, ("Discontinuous Lagrange", 0),)
    domain.topology.create_entities(tdim)
    domain.topology.create_entities(tdim-1)
    domain.topology.create_connectivity(tdim,tdim)
    cell_map = domain.topology.index_map(tdim)

    # Parameters
    dt = fem.Constant(domain, 0.1)
    k   = fem.Constant(domain, 1.0)
    c_p = fem.Constant(domain, 1.0)
    rho = fem.Constant(domain, 1.0)
    T_hot = 2.0
    T_cold = 0.0
    num_timesteps = 5

    unp1 = fem.Function(CG1, name="unp1")
    un   = fem.Function(CG1, name="un")
    udc  = fem.Function(CG1, name="dirichlet")
    du   = fem.Function(CG1, name="delta_uh")
    alg_rhs = fem.Function(CG1, name="rhs")

    # Dirichlet BC
    def left_marker_dirichlet(x):
        return np.logical_or( np.isclose(x[1],1), np.logical_or(
                np.isclose(x[0],0), np.isclose(x[1],0)) )
    dirichlet_dofs = fem.locate_dofs_geometrical(CG1,left_marker_dirichlet)
    dirichlet_bcs = [fem.dirichletbc(udc, dirichlet_dofs)]
    
    # Forms
    dx = ufl.Measure("dx")

    # Incremental solve
    (u, v) = (unp1, ufl.TestFunction(CG1))
    a_ufl = k*ufl.dot(ufl.grad(u), ufl.grad(v))*dx
    a_ufl += (rho*c_p/dt)*u*v*dx
    l_ufl = (rho*c_p/dt)*un*v*dx
    r_ufl = a_ufl - l_ufl

    j_ufl = ufl.derivative(a_ufl,unp1)
    j_compiled = fem.form(j_ufl)
    r_compiled = fem.form(r_ufl)
    A = dolfinx.fem.petsc.create_matrix(j_compiled)
    L = dolfinx.fem.petsc.create_vector(r_compiled)
    x = dolfinx.fem.petsc.create_vector(r_compiled)

    time = 0.0
    unp1.x.array[:] = T_cold
    unp1.x.array[dirichlet_dofs] = T_hot
    udc.x.array[:] = T_hot
    vtk_writer = io.VTKFile(domain.comm, f"post/incremenetal.pvd", "wb")
    for titer in range(num_timesteps):
        time += dt.value
        un.x.array[:] = unp1.x.array[:]
        # ASSEMBLE
        A.zeroEntries()
        dolfinx.fem.petsc.assemble_matrix(A, j_compiled, bcs=dirichlet_bcs,)
        A.assemble()
        with L.localForm() as l_local:
            l_local.set(0.0)
        dolfinx.fem.petsc.assemble_vector(L,
                                          r_compiled,
                                          )
        L.ghostUpdate(addv=PETSc.InsertMode.ADD_VALUES, mode=PETSc.ScatterMode.REVERSE)
        L.scale(-1)

        # Dirichlet
        dolfinx.fem.petsc.apply_lifting(L, [j_compiled], [dirichlet_bcs], x0=[unp1.x.petsc_vec])
        dolfinx.fem.petsc.set_bc(L,dirichlet_bcs, x0=unp1.x.petsc_vec)
        L.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)

        # RHS to function
        alg_rhs.x.array[:] = L.array[:]

        # LINEAR SOLVE
        opts = {"pc_type" : "lu", "pc_factor_mat_solver_type" : "mumps",}
        x = du.x.petsc_vec
        with x.localForm() as x_local:
            x_local.set(0.0)
        ksp = petsc4py.PETSc.KSP()
        ksp.create(domain.comm)
        ksp.setOperators(A)
        ksp_opts = PETSc.Options()
        for k,v in opts.items():
            ksp_opts[k] = v
        ksp.setFromOptions()
        ksp.solve(L, x)
        x.ghostUpdate(addv=petsc4py.PETSc.InsertMode.INSERT, mode=petsc4py.PETSc.ScatterMode.FORWARD)
        ksp.destroy()
        du.x.scatter_forward()
        unp1.x.array[:] += du.x.array[:]

        vtk_writer.write_function([unp1,un,udc,alg_rhs],t=time)
    vtk_writer.close()
    for la_ds in [A, L, x]:
        la_ds.destroy()

if __name__=="__main__":
    main()

[francesco-ballarin]

@ordinary-slim

        multiphenicsx.fem.petsc.set_bc(L,dirichlet_bcs, x0=unp1.x.petsc_vec, restriction=restriction)

This one is odd. The typical pattern is to pass in the solutions themselves, so I would say x0=[unp1].
I can have another look if this is not enough to fix the issue.

Please keep in mind that that call to set_bc has changed in dolfinx 0.9, and will need to change correspondingly in multiphenicsx once you upgrade both.

[ordinary-slim]

It doesn’t run if I pass x0=unp1 or x0=[unp1]. I get a TypeError: __init__(): incompatible function arguments. Note that I am on the latest commits of multiphenicsx and dolfinx so the typing of set_bc does say x0: typing.Optional[petsc4py.PETSc.Vec] = None. Could you have another look please?

EDIT: the multiphenicsx approach works if I avoid the x0 argument by directly modifying the Dirichlet condition to u_{dc} \longleftarrow u_{dc} - u^{n+1}. I think this might be a bug in multiphenicsx.fem.petsc.set_bc’s x0 argument.

[francesco-ballarin]

I will try having a look, but it will take me a couple of days.

I have a test case multiphenicsx/tests/unit/fem/test_assembler_restriction.py at dolfinx-v0.8.0 · multiphenics/multiphenicsx · GitHub (for dolfinx 0.8.0, change the branch name if your version is different) that covers that, so I don’t expect that there is a fundamental mistake in the implementation. Can you spot any big difference between the steps that are done in the test and your code?

[ordinary-slim]

Your link did make me realize that apply_lifting can take an optional argument restriction_x0. It is not the case for set_bc and set_bc expects a restricted x0 but I was passing an unrestricted vector to it. I think that’s the issue.

EDIT: In fact your x0 in your test_vector_assembly_with_restriction is your restricted solution

    restricted_fem_module.petsc.set_bc(
        restricted_vector_nonlinear, bcs, restricted_solution, restriction=dofmap_restriction)

[ordinary-slim]

@francesco-ballarin could you consider modifying set_bc to expect an unrestricted x0 by default? I did it like this:

def set_bc(  # type: ignore[no-any-unimported]
    b: petsc4py.PETSc.Vec, bcs: list[dolfinx.fem.DirichletBC] = [],
    x0: typing.Optional[petsc4py.PETSc.Vec] = None,
    alpha: float = 1.0,
    restriction: typing.Optional[mcpp.fem.DofMapRestriction] = None,
    restriction_x0: typing.Optional[mcpp.fem.DofMapRestriction] = None,
    function_space_x0: typing.Optional[dolfinx.fem.FunctionSpace] = None,
) -> None:
    """Apply the function :func:`dolfinx.fem.set_bc` to a PETSc Vector."""
    if restriction is None:
        if x0 is not None:
            x0 = x0.array_r
        for bc in bcs:
            bc.set(b.array_w, x0, alpha)
    else:
        x0_wrapper = x0
        if restriction_x0 is not None:
            if function_space_x0 is None:
                raise ValueError("restriction_x0 and function_space_x0 must \
                be provided together.")
            dofmap_x0 = function_space_x0.dofmap
            x0_wrapper = VecSubVectorReadWrapper(x0, dofmap_x0, restriction_x0, ghosted=False)
        with VecSubVectorWrapper(b, restriction.dofmap, restriction, ghosted=False) as b_sub, \
                x0_wrapper as x0_sub:
            for bc in bcs:
                bc.set(b_sub, x0_sub, alpha)

but I expect your version will be much cleaner. I think this default behavior makes more sense.

[francesco-ballarin]

I agree, it would probably make more sense, but having function_space_x0 is a deal breaker for me. I guess it’s not really needed (in the sense that restriction_x0 will probably already have the dofmap of function_space_x0), but I can try to look into it.

This is a very busy period for me, so it would be beneficial if I could get some help.
For instance:

1. would you be willing to open a PR (targeting the main branch, i.e. after 0.9.0, I don’t have the manpower to backport fixes to old releases) that adds the documentation to every function in multiphenicsx/fem/petsc.py, especially the ones that you mentioned (setting BCs and applying lifting). If the interface is different between the two (e.g., one time x0 is restricted, the other isn’t) that should be noted explicitly, because that is the relevant bit of information that you were missing (and I forgot, otherwise I would have told you from the very first post ). If you look in the file there are several examples like

    """
    ... a short description (probably already written) ...

    Parameters
    ----------
    ....

Please document the functions as they are now, if/when we change them with your proposal we’ll update the documentation later on.

2. can you help me track down where these functions about set_bc and lifting are used in the tutorials? It’s possible that some of them are only used in tests, and that may have been part of the challenge you encountered.