import petsc4py
# petsc4py.init(['-info'])
from petsc4py.PETSc import KSP, IS, Mat, Options, Viewer, PC, Vec, NullSpace, MatPartitioning, COMM_SELF


NEQ = 6

# Simple parallel diagonal matrix
M = Mat().create()
M.setSizes([NEQ, NEQ])
M.setType('aij')
M.setUp()
M.setPreallocationNNZ(1)
cs, ce = M.getOwnershipRange()
for row in range(cs, ce):
    M.setValue(row, row, row)
M.assemble()
M.view()


ovl=2
# reproduce initial layout
isovl = IS().createStride(ce-cs, cs, 1, comm=COMM_SELF)
loc = isovl.getLocalSize()
print(f"locSize before = {loc}", flush=True)
isovl.view()

# increase overlap
M.increaseOverlap(isovl, ovl)
loc = isovl.getLocalSize()
print(f"locSize after = {loc}", flush=True)
isovl.view()