Very rough but works.

doc/demo/demo_plasticity_von_mises.py

@@ -414,38 +414,58 @@ def sigma_external(derivatives):
 #
 # At the same time, we estimate the compilation overhead caused by the first call
 # of JIT-ed Numba functions.
-
-# %%
-# We need to initialize `Du` with small values in order to avoid the division by
-# zero error
-eps = np.finfo(PETSc.ScalarType).eps
-Du.x.array[:] = eps
-
-timer = common.Timer("DOLFINx_timer")
-timer.start()
-evaluated_operands = evaluate_operands(F_external_operators)
-_ = evaluate_external_operators(J_external_operators, evaluated_operands)
-timer.stop()
-pass_1 = timer.elapsed()[0]
-
-timer.start()
-evaluated_operands = evaluate_operands(F_external_operators)
-_ = evaluate_external_operators(J_external_operators, evaluated_operands)
-timer.stop()
-pass_2 = timer.elapsed()[0]
-
-print(f"\nNumba's JIT compilation overhead: {pass_1 - pass_2}")
-
-
-# %% [markdown]
+#
 # ### Solving the problem
 #
-# Solving the problem is carried out in a manually implemented Newton solver.
-
 # %%
 u = fem.Function(V, name="displacement")
-du = fem.Function(V, name="Newton_correction")
-external_operator_problem = LinearProblem(J_replaced, F_replaced, Du, bcs=bcs)
+from dolfinx.fem.petsc import NonlinearProblem, assemble_vector, assemble_matrix_mat, apply_lifting, set_bc
+from petsc4py import PETSc
+
+class PlasticityProblem(NonlinearProblem):
+    def form(self, x: PETSc.Vec) -> None:
+        """This function is called before the residual or Jacobian is
+        computed. This is usually used to update ghost values.
+
+        Args:
+           x: The vector containing the latest solution
+        """
+        x.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)
+
+        evaluated_operands = evaluate_operands(F_external_operators) 
+        ((_, sigma_new, dp_new),) = evaluate_external_operators(J_external_operators, evaluated_operands)
+
+        sigma.ref_coefficient.x.array[:] = sigma_new
+        dp.x.array[:] = dp_new
+
+    def F(self, x: PETSc.Vec, b: PETSc.Vec) -> None:
+        """Assemble the residual F into the vector b.
+
+        Args:
+            x: The vector containing the latest solution
+            b: Vector to assemble the residual into
+        """
+        # Reset the residual vector
+        with b.localForm() as b_local:
+            b_local.set(0.0)
+        assemble_vector(b, self._L)
+
+        # Apply boundary condition
+        apply_lifting(b, [self._a], bcs=[self.bcs], x0=[x], alpha=-1.0)
+        b.ghostUpdate(addv=PETSc.InsertMode.ADD, mode=PETSc.ScatterMode.REVERSE)
+        set_bc(b, self.bcs, x, -1.0)
+
+    def J(self, x: PETSc.Vec, A: PETSc.Mat) -> None:
+        """Assemble the Jacobian matrix.
+
+        Args:
+            x: The vector containing the latest solution
+        """
+        A.zeroEntries()
+        assemble_matrix_mat(A, self._a, self.bcs)
+        A.assemble()
+
+problem = PlasticityProblem(F_replaced, Du, bcs=bcs, J=J_replaced)
 
 # %%
 # Defining a cell containing (Ri, 0) point, where we calculate a value of u
@@ -462,44 +482,24 @@ def sigma_external(derivatives):
 loadings = q_lim * load_steps
 results = np.zeros((num_increments, 2))
 
-for i, loading_v in enumerate(loadings):
-    loading.value = loading_v
-    external_operator_problem.assemble_vector()
-
-    residual_0 = residual = external_operator_problem.b.norm()
-    Du.x.array[:] = 0.0
-
-    if MPI.COMM_WORLD.rank == 0:
-        print(f"\nresidual , {residual} \n increment: {i+1!s}, load = {loading.value}")
-
-    for iteration in range(0, max_iterations):
-        if residual / residual_0 < relative_tolerance:
-            break
-        external_operator_problem.assemble_matrix()
-        external_operator_problem.solve(du)
-        du.x.scatter_forward()
-
-        Du.x.petsc_vec.axpy(-1.0, du.x.petsc_vec)
-        Du.x.scatter_forward()
+from dolfinx.nls.petsc import NewtonSolver
+solver = NewtonSolver(mesh.comm, problem)
 
-        evaluated_operands = evaluate_operands(F_external_operators)
+ksp = solver.krylov_solver
+opts = PETSc.Options()  # type: ignore
+option_prefix = ksp.getOptionsPrefix()
+opts[f"{option_prefix}ksp_type"] = "preonly"
+opts[f"{option_prefix}pc_type"] = "lu"
+opts[f"{option_prefix}pc_factor_mat_solver_type"] = "mumps"
+ksp.setFromOptions()
 
-        # Implementation of an external operator may return several outputs and
-        # not only its evaluation. For example, `C_tang_impl` returns a tuple of
-        # Numpy-arrays with values of `C_tang`, `sigma` and `dp`.
-        ((_, sigma_new, dp_new),) = evaluate_external_operators(J_external_operators, evaluated_operands)
-
-        # In order to update the values of the external operator we may directly
-        # access them and avoid the call of
-        # `evaluate_external_operators(F_external_operators, evaluated_operands).`
-        sigma.ref_coefficient.x.array[:] = sigma_new
-        dp.x.array[:] = dp_new
+eps = np.finfo(PETSc.ScalarType).eps
 
-        external_operator_problem.assemble_vector()
-        residual = external_operator_problem.b.norm()
+for i, loading_v in enumerate(loadings):
+    loading.value = loading_v
+    Du.x.array[:] = eps
 
-        if MPI.COMM_WORLD.rank == 0:
-            print(f"    it# {iteration} residual: {residual}")
+    solver.solve(Du)
 
     u.x.petsc_vec.axpy(-1.0, Du.x.petsc_vec)
     u.x.scatter_forward()