Parallel code for extraction of submatrices

Hello,
I am running in parallel for the solution of some PDE. After this, I have on every process a matrix, which I would like to partition accordingly for further use. See the code below to understand better.

from petsc4py import PETSc
import dolfinx, ufl
from mpi4py import MPI

from ufl import inner, grad, dx

import numpy as np

comm = MPI.COMM_WORLD

Lx, Ly = 1, 1
px, py = 2, 2
N = 50

assert comm.size == px * py

dim = 2

mesh, rank_bounds, primal_indices = partition_domain(comm, px, py, Lx, Ly, N, dim)

V = dolfinx.fem.VectorFunctionSpace(mesh, ('CG', 1))
u, v = ufl.TrialFunction(V), ufl.TestFunction(V)

E = 1
nu = 0.3
mu = E / (2.0 * (1.0 + nu))
lmbda = E * nu / ((1.0 + nu) * (1.0 - 2.0 * nu))

def epsilon(v):
    return ufl.sym(grad(v))

def sigma(v):
    return 2.0 * mu * epsilon(v) + lmbda * ufl.tr(epsilon(v)) * ufl.Identity(len(v))

boundaries = [(1, lambda x: np.isclose(x[0], 0)),
                       (2, lambda x: np.isclose(x[0], px*Lx))]

facet_indices, facet_markers = [], []
for (marker, locator) in boundaries:
      facets = dolfinx.mesh.locate_entities(mesh, mesh.topology.dim-1, locator)
      facet_indices.append(facets)
      facet_markers.append(np.full_like(facets, marker))
      facet_indices = np.hstack(facet_indices).astype(np.int32)
      facet_markers = np.hstack(facet_markers).astype(np.int32)
      sorted_facets = np.argsort(facet_indices)

facet_tag = dolfinx.mesh.meshtags(mesh, mesh.topology.dim-1, facet_indices[sorted_facets], facet_markers[sorted_facets])
ds = ufl.ds(subdomain_data=facet_tag, domain=mesh)
t = ufl.as_vector([0.0, -1.0])
a = inner(sigma(u), grad(v))
L = inner(t, v)

u_D = dolfinx.fem.Function(V)
facets = facet_tag.find(1)
clamped_edge = dolfinx.fem.locate_dofs_topological(V, mesh.topology.dim-1, facets)
zero_bc = dolfinx.fem.dirichletbc(np.array((0,) * mesh.geometry.dim, dtype=np.double), clamped_edge, V)
A = dolfinx.fem.petsc.assemble_matrix(dolfinx.fem.form(a * dx), bcs=[zero_bc])
A.assemble()

## split dofs into some partition

A_ = PETSc.Mat().createAIJ(#some size#, comm=PETSc.COMM_SELF).setUp()
A_.setValues(range(A_.getSizes()[0][0]), range(A_.getSizes()[1][0]), A.getValues(.., ..))

However, the code below remains stuck on the setting of the values for the submatrix Arr. Any help is appreciated. If I stick to the global communicator then the code goes through but throws mistakes because of wrong sizes.

In the MPI.COMM_SELF case, I get the following error

Loguru caught a signal: SIGSEGV


Loguru caught a signal: SIGSEGV
Loguru caught a signal: SIGSEGV

Loguru caught a signal: SIGSEGV
Stack trace:
22      0x564b2db39245 _start + 37
21      0x7f4238b6ce40 __libc_start_main + 128
20      0x7f4238b6cd90 /lib/x86_64-linux-gnu/libc.so.6(+0x29d90) [0x7f4238b6cd90]
19      0x564b2db3934d Py_BytesMain + 45
18      0x564b2db630ae Py_RunMain + 702
17      0x564b2db71e33 _PyRun_AnyFileObject + 67
16      0x564b2db72138 _PyRun_SimpleFileObject + 424
15      0x564b2db72c55 python3(+0x265c55) [0x564b2db72c55]
14      0x564b2db6bd5b python3(+0x25ed5b) [0x564b2db6bd5b]
13      0x564b2db72f08 python3(+0x265f08) [0x564b2db72f08]
12      0x564b2db46456 PyEval_EvalCode + 134
11      0x564b2da4e766 python3(+0x141766) [0x564b2da4e766]
10      0x564b2da57e39 _PyEval_EvalFrameDefault + 25849
9       0x564b2da5f77c _PyObject_MakeTpCall + 508
8       0x564b2da73639 python3(+0x166639) [0x564b2da73639]
7       0x564b2da5e964 _PyObject_FastCallDictTstate + 196
6       0x564b2da5214a _PyEval_EvalFrameDefault + 2058
5       0x564b2da6a269 python3(+0x15d269) [0x564b2da6a269]
4       0x7f423861ce82 /usr/local/lib/python3.10/dist-packages/petsc4py/lib/linux-gnu-real-32/PETSc.cpython-310-x86_64-linux-gnu.so(+0x244e82) [0x7f423861ce82]
3       0x7f423861b207 /usr/local/lib/python3.10/dist-packages/petsc4py/lib/linux-gnu-real-32/PETSc.cpython-310-x86_64-linux-gnu.so(+0x243207) [0x7f423861b207]
2       0x7f4237468001 MatSetValues + 241
1       0x7f42375b35e4 /usr/local/petsc/linux-gnu-real-32/lib/libpetsc.so.3.17(+0x9195e4) [0x7f42375b35e4]
0       0x7f4238b85520 /lib/x86_64-linux-gnu/libc.so.6(+0x42520) [0x7f4238b85520]
2023-11-16 15:48:38.743 (   2.008s) [main            ]                       :0     FATL| Signal: SIGSEGV
Stack trace:
22      0x560c842b7245 _start + 37
21      0x7f66a59f3e40 __libc_start_main + 128
20      0x7f66a59f3d90 /lib/x86_64-linux-gnu/libc.so.6(+0x29d90) [0x7f66a59f3d90]
19      0x560c842b734d Py_BytesMain + 45
18      0x560c842e10ae Py_RunMain + 702
17      0x560c842efe33 _PyRun_AnyFileObject + 67
16      0x560c842f0138 _PyRun_SimpleFileObject + 424
15      0x560c842f0c55 python3(+0x265c55) [0x560c842f0c55]
14      0x560c842e9d5b python3(+0x25ed5b) [0x560c842e9d5b]
13      0x560c842f0f08 python3(+0x265f08) [0x560c842f0f08]
12      0x560c842c4456 PyEval_EvalCode + 134
11      0x560c841cc766 python3(+0x141766) [0x560c841cc766]
10      0x560c841d5e39 _PyEval_EvalFrameDefault + 25849
9       0x560c841dd77c _PyObject_MakeTpCall + 508
8       0x560c841f1639 python3(+0x166639) [0x560c841f1639]
7       0x560c841dc964 _PyObject_FastCallDictTstate + 196
6       0x560c841d014a _PyEval_EvalFrameDefault + 2058
5       0x560c841e8269 python3(+0x15d269) [0x560c841e8269]
4       0x7f66a54a3e82 /usr/local/lib/python3.10/dist-packages/petsc4py/lib/linux-gnu-real-32/PETSc.cpython-310-x86_64-linux-gnu.so(+0x244e82) [0x7f66a54a3e82]
3       0x7f66a54a2207 /usr/local/lib/python3.10/dist-packages/petsc4py/lib/linux-gnu-real-32/PETSc.cpython-310-x86_64-linux-gnu.so(+0x243207) [0x7f66a54a2207]
2       0x7f66a42ef001 MatSetValues + 241
1       0x7f66a443a5e4 /usr/local/petsc/linux-gnu-real-32/lib/libpetsc.so.3.17(+0x9195e4) [0x7f66a443a5e4]
0       0x7f66a5a0c520 /lib/x86_64-linux-gnu/libc.so.6(+0x42520) [0x7f66a5a0c520]
2023-11-16 15:48:38.743 (   2.015s) [main            ]                       :0     FATL| Signal: SIGSEGV
Stack trace:
22      0x55b79c03d245 _start + 37
21      0x7f396a6b7e40 __libc_start_main + 128
20      0x7f396a6b7d90 /lib/x86_64-linux-gnu/libc.so.6(+0x29d90) [0x7f396a6b7d90]
19      0x55b79c03d34d Py_BytesMain + 45
18      0x55b79c0670ae Py_RunMain + 702
17      0x55b79c075e33 _PyRun_AnyFileObject + 67
16      0x55b79c076138 _PyRun_SimpleFileObject + 424
15      0x55b79c076c55 python3(+0x265c55) [0x55b79c076c55]
14      0x55b79c06fd5b python3(+0x25ed5b) [0x55b79c06fd5b]
13      0x55b79c076f08 python3(+0x265f08) [0x55b79c076f08]
12      0x55b79c04a456 PyEval_EvalCode + 134
11      0x55b79bf52766 python3(+0x141766) [0x55b79bf52766]
10      0x55b79bf5be39 _PyEval_EvalFrameDefault + 25849
9       0x55b79bf6377c _PyObject_MakeTpCall + 508
8       0x55b79bf77639 python3(+0x166639) [0x55b79bf77639]
7       0x55b79bf62964 _PyObject_FastCallDictTstate + 196
6       0x55b79bf5614a _PyEval_EvalFrameDefault + 2058
5       0x55b79bf6e269 python3(+0x15d269) [0x55b79bf6e269]
4       0x7f396a167e82 /usr/local/lib/python3.10/dist-packages/petsc4py/lib/linux-gnu-real-32/PETSc.cpython-310-x86_64-linux-gnu.so(+0x244e82) [0x7f396a167e82]
3       0x7f396a166207 /usr/local/lib/python3.10/dist-packages/petsc4py/lib/linux-gnu-real-32/PETSc.cpython-310-x86_64-linux-gnu.so(+0x243207) [0x7f396a166207]
2       0x7f3968fb3001 MatSetValues + 241
1       0x7f39690fe5e4 /usr/local/petsc/linux-gnu-real-32/lib/libpetsc.so.3.17(+0x9195e4) [0x7f39690fe5e4]
0       0x7f396a6d0520 /lib/x86_64-linux-gnu/libc.so.6(+0x42520) [0x7f396a6d0520]
2023-11-16 15:48:38.744 (   2.030s) [main            ]                       :0     FATL| Signal: SIGSEGV
Stack trace:
22      0x55563f31a245 _start + 37
21      0x7f50d9d9fe40 __libc_start_main + 128
20      0x7f50d9d9fd90 /lib/x86_64-linux-gnu/libc.so.6(+0x29d90) [0x7f50d9d9fd90]
19      0x55563f31a34d Py_BytesMain + 45
18      0x55563f3440ae Py_RunMain + 702
17      0x55563f352e33 _PyRun_AnyFileObject + 67
16      0x55563f353138 _PyRun_SimpleFileObject + 424
15      0x55563f353c55 python3(+0x265c55) [0x55563f353c55]
14      0x55563f34cd5b python3(+0x25ed5b) [0x55563f34cd5b]
13      0x55563f353f08 python3(+0x265f08) [0x55563f353f08]
12      0x55563f327456 PyEval_EvalCode + 134
11      0x55563f22f766 python3(+0x141766) [0x55563f22f766]
10      0x55563f238e39 _PyEval_EvalFrameDefault + 25849
9       0x55563f24077c _PyObject_MakeTpCall + 508
8       0x55563f254639 python3(+0x166639) [0x55563f254639]
7       0x55563f23f964 _PyObject_FastCallDictTstate + 196
6       0x55563f23314a _PyEval_EvalFrameDefault + 2058
5       0x55563f24b269 python3(+0x15d269) [0x55563f24b269]
4       0x7f50d984fe82 /usr/local/lib/python3.10/dist-packages/petsc4py/lib/linux-gnu-real-32/PETSc.cpython-310-x86_64-linux-gnu.so(+0x244e82) [0x7f50d984fe82]
3       0x7f50d984e207 /usr/local/lib/python3.10/dist-packages/petsc4py/lib/linux-gnu-real-32/PETSc.cpython-310-x86_64-linux-gnu.so(+0x243207) [0x7f50d984e207]
2       0x7f50d869b001 MatSetValues + 241
1       0x7f50d87e65e4 /usr/local/petsc/linux-gnu-real-32/lib/libpetsc.so.3.17(+0x9195e4) [0x7f50d87e65e4]
0       0x7f50d9db8520 /lib/x86_64-linux-gnu/libc.so.6(+0x42520) [0x7f50d9db8520]
2023-11-16 15:48:38.743 (   2.023s) [main            ]                       :0     FATL| Signal: SIGSEGV

===================================================================================
=   BAD TERMINATION OF ONE OF YOUR APPLICATION PROCESSES
=   PID 138 RUNNING AT b0b84aa4a907
=   EXIT CODE: 11
=   CLEANING UP REMAINING PROCESSES
=   YOU CAN IGNORE THE BELOW CLEANUP MESSAGES
===================================================================================
YOUR APPLICATION TERMINATED WITH THE EXIT STRING: Segmentation fault (signal 11)
This typically refers to a problem with your application.
Please see the FAQ page for debugging suggestions