Issues with translating code from dolfin to dolfinx

Hello, I have the following code in dolfin, that I am trying to translate to dolfinx.

    up_sol.set_operator(A)

    while rel_res > rtol and residual > atol and iter < max_it:

        A = assemble(J_nonlinear,
                     form_compiler_parameters=compiler_parameters,
                     keep_diagonal=True
        )
        A.ident_zeros()
        [bc.apply(A) for bc in bcs]

        b = assemble(-F)

        [bc.apply(b, dvp_["n"].vector()) for bc in bcs]
        up_sol.solve(dvp_res.vector(), b)
        dvp_["n"].vector().axpy(lmbda, dvp_res.vector())
        [bc.apply(dvp_["n"].vector()) for bc in bcs]

For the moment, I am doing something like this, but I get different results and I am trying to understand whether this part of the code is the culprit.

    dvp_res = dvp_["n"].vector.copy()
    dvp_res.zeroEntries()

    ksp = PETSc.KSP().create(comm=MPI.COMM_WORLD)
    ksp.setType("preonly")
    ksp.getPC().setType("lu")
    ksp.getPC().setFactorSolverType("mumps")

    while rel_res > rtol and residual > atol and iter < max_it:
        chi = ufl.TrialFunction(dvp_["n"].function_space)
        J = ufl.derivative(F, dvp_["n"], chi)
        A = dolfinx.fem.petsc.assemble_matrix(dolfinx.fem.form(J))
        A.assemble()

        zero_rows = A.findZeroRows().array
        diagonal_values = A.getDiagonal().array
        diagonal_values[zero_rows] = 1.0
        diagonal_petsc = PETSc.Vec().createWithArray(diagonal_values)

        A.setOption(PETSc.Mat.Option.NEW_NONZERO_ALLOCATION_ERR, False)
        A.setDiagonal(diagonal_petsc)
        diagonal_petsc.destroy()
        A.assemble()

        for bc in bcs:
            dofs, _ = bc.dof_indices()
            A.zeroRowsColumns(dofs, diag=1)

        ksp.setOperators(A)

        b = dolfinx.fem.petsc.assemble_vector(dolfinx.fem.form(-F))
        b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
        dolfinx.fem.petsc.set_bc(b, bcs)
        b.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)

        ksp.solve(b, dvp_res)
       
        dvp_["n"].vector.axpy(lmbda, dvp_res)
        dvp_["n"].vector.ghostUpdate(addv=PETSc.InsertMode.ADD,      mode=PETSc.ScatterMode.REVERSE)
        dolfinx.fem.petsc.set_bc(dvp_["n"].vector, bcs)
        dvp_["n"].vector.ghostUpdate(addv=PETSc.InsertMode.INSERT_VALUES, mode=PETSc.ScatterMode.FORWARD)

Is the translation of code corresponding to the setting of the boundary conditions and the .ident_zeros() call correct, or is there a clear mistake? Thank you in advance.

I don’t see any obvious errors.
However, what is your use-case for ident_zeros? There might be better ways of performing the task you have in mind than modifying the PETSc matrix post assembly.

For instance: Modify matrix diagonal -- dolfinx version for A.ident_zeros() - #3 by dokken
illustrates this.
You might also be able to use dolfinx.mesh.create_submesh to restrict the function space instead.

Hi @dokken , thank you very much for your reply. You may have understood from the turtleFSI post that my use case is the Turek Hron FSI benchmark. I shared with you the code over there, let me know if you want me to attach it here as well for common visibility or if it is fine there.