Example for PETSc MatCreateLRC using FEniCS and petsc4py

poission_LRC.py

"""
Solve a simple Poisson system with a low rank update using PETSc MatCreateLRC.
The low rank update is defined at 3 boundaries with different weights, Dirichlet
BC is applied to the other boundary.

https://www.mcs.anl.gov/petsc/petsc-3.13/docs/manualpages/Mat/MatCreateLRC.html

some hints:
https://fenicsproject.discourse.group/t/add-matrix-to-assembled-matrix/122/2
https://lists.mcs.anl.gov/pipermail/petsc-users/2020-February/040228.html
https://fenicsproject.org/qa/11647/generating-matrices-applying-boundary-conditions-parallel/
"""
import matplotlib.pyplot as plt
import numpy as np
from dolfin import *
from petsc4py import PETSc

N = 16
mesh = UnitSquareMesh(N, N)

V = FunctionSpace(mesh, "P", 1)

p = TrialFunction(V)
q = TestFunction(V)

a = inner(grad(p), grad(q)) * dx
L = Constant(1) * q * dx

A = assemble(a)
b = assemble(L)

dbc = DirichletBC(V, Constant(0), "on_boundary && near(x[0], 0)")
dbc.apply(A, b)

bnd = [
    CompiledSubDomain("on_boundary && near(x[0], 1)"),
    CompiledSubDomain("on_boundary && near(x[1], 0)"),
    CompiledSubDomain("on_boundary && near(x[1], 1)"),
]
boundaries = MeshFunction("size_t", mesh, mesh.topology().dim() - 1, 0)
for i, bi in enumerate(bnd):
    bi.mark(boundaries, i)

ds = Measure("ds", domain=mesh, subdomain_data=boundaries)

delta = np.array([5.**k for k in range(len(bnd))])
f_lst = [Constant(d) * p * ds(i) for i, d in enumerate(delta)]
u_lst = [assemble(f) for f in f_lst]

# c = interpolate(Expression("N*x[0]", degree=1, N=N), V)
# c = interpolate(Constant(1), V)

nrows = A.size(0)
info(f"{nrows = }")

dofmap = V.dofmap()
# Local, global
sizes = [
    dofmap.index_map().size(IndexMap.MapSize.OWNED),
    dofmap.index_map().size(IndexMap.MapSize.GLOBAL),
]
info(f"{sizes = }")
comm = mesh.mpi_comm()

array = np.hstack([u.get_local() for u in u_lst])
ncol = len(u_lst)

D = PETSc.Mat().createDense(
    size=(sizes, (ncol, ncol)), array=array, comm=comm
)
D.setUp()
D.assemble()

# D_indpt, D_col, D_val = D.getValuesCSR()
# info(f"{D_indpt=}")
# info(f"{D_col=}")
# info(f"{D_val=}")

Ap = as_backend_type(A).mat()

use_diag = True
if use_diag:
    vec = PETSc.Vec().createSeq(size=len(delta))
    vec.setArray(delta)
    vec.assemble()

    info(f"{vec.getType()}")
    info(f"{delta}")
    info(f"{vec.array}")

    M = PETSc.Mat().createLRC(Ap, D, vec, D)

else:
    M = PETSc.Mat().createLRC(Ap, D, None, D)

ksp = PETSc.KSP().create()
PETScOptions.set("ksp_type", "gmres")
PETScOptions.set("pc_type", "gamg")
# XXX Matrix is not explicitly formed, LU not working!
# PETScOptions.set("ksp_type", "preonly")
# PETScOptions.set("pc_type", "lu")
PETScOptions.set("ksp_monitor_true_residual")
PETScOptions.set("ksp_converged_reason")
ksp.setFromOptions()

ksp.setOperators(M, Ap)
# ksp.setOperators(M)
ksp.setUp()

w = Function(V)
x = as_backend_type(w.vector())
ksp.solve(as_backend_type(b).vec(), x.vec())
x.update_ghost_values()

info(f"{norm(w)}")
plot(w)
plt.show()