I believe this is a bug in PETSc (report filed at: NEST combined with vinewtonrsls and bounds yields Could not find index set (#1831) · Issues · PETSc / petsc · GitLab).
Consider this hand-coded 1D finite element formulation of Poisson:
# Hand-coded formulation of the Poisson equation with homogeneous Dirichlet boundary conditions
# using PETSc
# Author: Jørgen S. Dokken
# SPDX-License-Identifier: MIT
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
# Define source term and boundary condition
def f(x_ref, a, b):
return np.pi**2 * np.sin(np.pi * (a + x_ref * (b - a)))
def g(x_ref, a, b):
# NOTE: The implementation below only works for homogeneous BCs
return 0.0
def phi_0(x_ref):
return 1 - x_ref
def phi_1(x_ref):
return x_ref
def dphi_0(x_ref, b, a):
return -1 / (b - a)
def dphi_1(x_ref, b, a):
return 1 / (b - a)
# Equal spacing option
# x_start = 0.0
# x_end = 1.0
# x = np.linspace(x_start, x_end, N, endpoint=True)
# N = 15
x = np.array([0.0, 0.05, 0.1, 0.3, 0.5, 0.6, 0.8, 1.0])
N = len(x)
# Create dofmap
dofmap = np.repeat(np.arange(N), 2).reshape(-1, 2).astype(np.int32)
dofmap[:, 1] += 1
# Create submatrices
A00 = PETSc.Mat().createAIJ([N, N])
A11 = PETSc.Mat().createAIJ([N, N])
A = PETSc.Mat().createNest([[A00, None], [None, A11]])
# Create subvectors
b0 = PETSc.Vec().createSeq(N)
b1 = PETSc.Vec().createSeq(N)
b = PETSc.Vec().createNest([b0, b1])
compute_interval_jacobian = lambda x_ref, a, b: b - a
def fill_matrix(mat):
"""
Compute stiffness matrix with a one-point quadrature
"""
mat.zeroEntries()
p_ref = 0.5
for cell in range(N - 1):
a = x[cell]
b = x[cell + 1]
J_e = compute_interval_jacobian(p_ref, a, b)
Ke = np.array(
[
[
dphi_0(p_ref, b, a) * dphi_0(p_ref, b, a) * J_e,
dphi_0(p_ref, b, a) * dphi_1(p_ref, b, a) * J_e,
],
[
dphi_1(p_ref, b, a) * dphi_0(p_ref, b, a) * J_e,
dphi_1(p_ref, b, a) * dphi_1(p_ref, b, a) * J_e,
],
],
dtype=np.float64,
)
if cell == 0:
Ke[0, :] = 0.0
Ke[0, 0] = 1.0
elif cell == N - 2:
Ke[1, :] = 0.0
Ke[1, 1] = 1.0
rows = dofmap[cell, :]
mat.setValues(rows, rows, Ke, addv=True)
mat.assemble()
def fill_vector(vec):
"""
Compute load vector with a one-point quadrature
"""
with vec.localForm() as local_vec:
local_vec.set(0.0)
for cell in range(N - 1):
a = x[cell]
b = x[cell + 1]
J_e = compute_interval_jacobian(0.0, a, b)
fe = np.array(
[
f(0.5, a, b) * phi_0(0.5) * J_e,
f(0.5, a, b) * phi_1(0.5) * J_e,
],
dtype=np.float64,
)
if cell == 0:
fe[0] = g(0.0, a, b)
elif cell == N - 2:
fe[1] = g(1.0, a, b)
rows = dofmap[cell, :]
vec.setValues(rows, fe, addv=True)
vec.assemble()
# Fill A and b as it only has to be done once
fill_matrix(A00)
fill_matrix(A11)
A.assemble()
fill_vector(b0)
fill_vector(b1)
b.assemble()
def assemble_jacobian(_snes, x, J, P):
"""Wrapper for snes"""
J.zeroEntries()
J00 = J.getNestSubMatrix(0, 0)
J11 = J.getNestSubMatrix(1, 1)
A00.copy(J00)
A11.copy(J11)
J.assemble()
# Working arrays for residual
tmp0 = b0.duplicate()
tmp1 = b1.duplicate()
def assemble_residual(_snes, x, F):
"""Wrapper for snes"""
F0, F1 = F.getNestSubVecs()
x0, x1 = x.getNestSubVecs()
with F0.localForm() as local_F0, F1.localForm() as local_F1:
local_F0.set(0.0)
local_F1.set(0.0)
b0.copy(F0)
b1.copy(F1)
F0.scale(-1)
F1.scale(-1)
with tmp0.localForm() as local_tmp0, tmp1.localForm() as local_tmp1:
local_tmp0.set(0.0)
local_tmp1.set(0.0)
with tmp0.localForm() as local_tmp0:
local_tmp0.set(0.0)
A00.mult(x0, tmp0)
A11.mult(x1, tmp1)
F0.axpy(1.0, tmp0)
F1.axpy(1.0, tmp1)
F0.array_w[0] = 0.0
F0.array_w[-1] = 0.0
F1.array_w[0] = 0.0
F1.array_w[-1] = 0.0
F.assemble()
snes = PETSc.SNES().create(MPI.COMM_SELF)
F = b.copy()
J = A.copy()
snes.setFunction(assemble_residual, F)
snes.setJacobian(assemble_jacobian, J, None)
solver_lower = b.copy()
for sub_vec in solver_lower.getNestSubVecs():
with sub_vec.localForm() as local_sub_vec:
local_sub_vec.set(-0.5)
solver_upper = b.copy()
for sub_vec in solver_upper.getNestSubVecs():
with sub_vec.localForm() as local_sub_vec:
local_sub_vec.set(0.5)
snes.setVariableBounds(solver_lower, solver_upper)
petsc_options = {
"snes_type": "vinewtonrsls",
# "snes_type": "newtonls",
"ksp_type": "preonly",
"pc_type": "lu",
"snes_linesearch_type": "none",
"pc_factor_mat_solver_type": "mumps",
"snes_monitor": None,
"snes_error_if_not_converged": True,
}
opts = PETSc.Options()
for key, val in petsc_options.items():
opts.setValue(key, val)
snes.setFromOptions()
nest_IS = J.getNestISs()
fieldsplit_IS = tuple([(f"{i}", IS) for i, IS in enumerate(nest_IS[0])])
snes.getKSP().getPC().setFieldSplitIS(*fieldsplit_IS)
snes.setErrorIfNotConverged(True)
uh = b.duplicate()
snes.solve(None, uh)
# Direct solver check
ksp = PETSc.KSP().create(MPI.COMM_SELF)
ksp.setOperators(A)
ksp.setType("preonly")
ksp.getPC().setType("lu")
ksp.getPC().setFactorSolverType("mumps")
ksp.setErrorIfNotConverged(True)
ksp.setMonitor(
lambda ksp, its, rnorm: print(f"KSP Iteration {its}, Residual norm: {rnorm}")
)
ksp.solve(b, uh)
import matplotlib.pyplot as plt
uh0, uh1 = uh.getNestSubVecs()
plt.plot(x, uh0.array, label="Numerical Solution (0)")
plt.plot(x, uh1.array, "bs", label="Numerical Solution (1)")
x_ref = np.linspace(0.0, 1.0, 100, endpoint=True)
plt.plot(x_ref, np.sin(np.pi * x_ref), label="Exact")
plt.legend()
plt.show()
This works with KSP, with SNES and “newtonls” and SNES with “vinewtonrsls” without bounds.
Adding the bounds yields:
0 SNES Function norm 3.904873750087e+00
1 SNES Function norm 2.902005876076e+00
Traceback (most recent call last):
File "/home/dokken/Documents/src/debug.py", line 224, in <module>
snes.solve(None, uh)
~~~~~~~~~~^^^^^^^^^^
File "petsc4py/PETSc/SNES.pyx", line 1738, in petsc4py.PETSc.SNES.solve
petsc4py.PETSc.Error: error code 75
[0] SNESSolve() at /home/conda/feedstock_root/build_artifacts/bld/rattler-build_petsc_1755537683/work/src/snes/interface/snes.c:4846
[0] SNESSolve_VINEWTONRSLS() at /home/conda/feedstock_root/build_artifacts/bld/rattler-build_petsc_1755537683/work/src/snes/impls/vi/rs/virs.c:381
[0] MatCreateSubMatrix() at /home/conda/feedstock_root/build_artifacts/bld/rattler-build_petsc_1755537683/work/src/mat/interface/matrix.c:8456
[0] MatCreateSubMatrix_Nest() at /home/conda/feedstock_root/build_artifacts/bld/rattler-build_petsc_1755537683/work/src/mat/impls/nest/matnest.c:712
[0] MatNestFindSubMat() at /home/conda/feedstock_root/build_artifacts/bld/rattler-build_petsc_1755537683/work/src/mat/impls/nest/matnest.c:699
[0] Arguments are incompatible
[0] Could not find index set