Isolate petsc4y in Python interface (#2703)

* Isolate PETSc imports * flake8 fixes * mypy fixes * flake8 fix * Newton fix * Fixes * Updates * Demo updates * flake8 fix * Various import fixes * Vector support * Updates * In progress * Updates * Updates * Doc fix * Update * Tidy up * Formatting * Minor demo improvements * Small edits * Simplify * PETSc improvements * Simplifications * Doc improvments * mypy fixes * Improve test imports * Test generalisation * Apply suggestions from code review Co-authored-by: Chris Richardson <chris@bpi.cam.ac.uk> * Fix small doc typos --------- Co-authored-by: Chris Richardson <chris@bpi.cam.ac.uk>
FEniCS · Jul 4, 2023 · 08006c4 · 08006c4
1 parent cb80654
commit 08006c4
Show file tree

Hide file tree

Showing 40 changed files with 721 additions and 579 deletions.
diff --git a/cpp/dolfinx/la/MatrixCSR.h b/cpp/dolfinx/la/MatrixCSR.h
@@ -357,7 +357,7 @@ class MatrixCSR
 
   /// Block size
   /// @return block sizes for rows and columns
-  const std::array<int, 2>& block_size() const { return _bs; }
+  std::array<int, 2> block_size() const { return _bs; }
 
 private:
   // Maps for the distribution of the ows and columns

diff --git a/cpp/dolfinx/la/Vector.h b/cpp/dolfinx/la/Vector.h
@@ -7,6 +7,7 @@
 #pragma once
 
 #include "utils.h"
+#include <cmath>
 #include <complex>
 #include <dolfinx/common/IndexMap.h>
 #include <dolfinx/common/Scatterer.h>
@@ -305,53 +306,60 @@ auto norm(const V& x, Norm type = Norm::l2)
 /// @param[in,out] basis The set of vectors to orthonormalise. The
 /// vectors must have identical parallel layouts. The vectors are
 /// modified in-place.
-/// @param[in] tol The tolerance used to detect a linear dependency
+/// @tparam V dolfinx::la::Vector
 template <class V>
-void orthonormalize(std::span<V> basis, double tol = 1.0e-10)
+void orthonormalize(std::vector<std::reference_wrapper<V>> basis)
 {
   using T = typename V::value_type;
+  using U = typename dolfinx::scalar_value_type_t<T>;
+
   // Loop over each vector in basis
   for (std::size_t i = 0; i < basis.size(); ++i)
   {
     // Orthogonalize vector i with respect to previously orthonormalized
     // vectors
+    V& bi = basis[i].get();
     for (std::size_t j = 0; j < i; ++j)
     {
+      const V& bj = basis[j].get();
+
       // basis_i <- basis_i - dot_ij  basis_j
-      T dot_ij = inner_product(basis[i], basis[j]);
-      std::transform(basis[j].array().begin(), basis[j].array().end(),
-                     basis[i].array().begin(), basis[i].mutable_array().begin(),
+      auto dot_ij = inner_product(bi, bj);
+      std::transform(bj.array().begin(), bj.array().end(), bi.array().begin(),
+                     bi.mutable_array().begin(),
                      [dot_ij](auto xj, auto xi) { return xi - dot_ij * xj; });
     }
 
     // Normalise basis function
-    double norm = la::norm(basis[i], la::Norm::l2);
-    if (norm < tol)
+    auto norm = la::norm(bi, la::Norm::l2);
+    if (norm * norm < std::numeric_limits<U>::epsilon())
     {
       throw std::runtime_error(
           "Linear dependency detected. Cannot orthogonalize.");
     }
-    std::transform(basis[i].array().begin(), basis[i].array().end(),
-                   basis[i].mutable_array().begin(),
+    std::transform(bi.array().begin(), bi.array().end(),
+                   bi.mutable_array().begin(),
                    [norm](auto x) { return x / norm; });
   }
 }
 
-/// Test if basis is orthonormal
-/// @param[in] basis The set of vectors to check
-/// @param[in] tol The tolerance used to test for orthonormality
-/// @return True is basis is orthonormal, otherwise false
+/// @brief Test if basis is orthonormal.
+/// @param[in] basis The set of vectors to check.
+/// @return True is basis is orthonormal, otherwise false.
 template <class V>
-bool is_orthonormal(std::span<const V> basis, double tol = 1.0e-10)
+bool is_orthonormal(std::vector<std::reference_wrapper<const V>> basis)
 {
   using T = typename V::value_type;
+  using U = typename dolfinx::scalar_value_type_t<T>;
+
+  // auto tol = std::sqrt(T(std::numeric_limits<U>::epsilon()));
   for (std::size_t i = 0; i < basis.size(); i++)
   {
     for (std::size_t j = i; j < basis.size(); j++)
     {
-      const double delta_ij = (i == j) ? 1.0 : 0.0;
-      T dot_ij = inner_product(basis[i], basis[j]);
-      if (std::abs(delta_ij - dot_ij) > tol)
+      T delta_ij = (i == j) ? T(1) : T(0);
+      auto dot_ij = inner_product(basis[i].get(), basis[j].get());
+      if (std::norm(delta_ij - dot_ij) > std::numeric_limits<U>::epsilon())
         return false;
     }
   }

diff --git a/cpp/dolfinx/la/petsc.cpp b/cpp/dolfinx/la/petsc.cpp
@@ -75,15 +75,24 @@ Vec la::petsc::create_vector(MPI_Comm comm, std::array<std::int64_t, 2> range,
 
   // Get local size
   assert(range[1] >= range[0]);
-  const std::int32_t local_size = range[1] - range[0];
+  std::int32_t local_size = range[1] - range[0];
 
   Vec x;
-  const std::vector<PetscInt> _ghosts(ghosts.begin(), ghosts.end());
-  ierr = VecCreateGhostBlock(comm, bs, bs * local_size, PETSC_DETERMINE,
-                             _ghosts.size(), _ghosts.data(), &x);
-  CHECK_ERROR("VecCreateGhostBlock");
-  assert(x);
+  std::vector<PetscInt> _ghosts(ghosts.begin(), ghosts.end());
+  if (bs == 1)
+  {
+    ierr = VecCreateGhost(comm, local_size, PETSC_DETERMINE, _ghosts.size(),
+                          _ghosts.data(), &x);
+    CHECK_ERROR("VecCreateGhost");
+  }
+  else
+  {
+    ierr = VecCreateGhostBlock(comm, bs, bs * local_size, PETSC_DETERMINE,
+                               _ghosts.size(), _ghosts.data(), &x);
+    CHECK_ERROR("VecCreateGhostBlock");
+  }
 
+  assert(x);
   return x;
 }
 //-----------------------------------------------------------------------------
@@ -94,8 +103,23 @@ Vec la::petsc::create_vector_wrap(const common::IndexMap& map, int bs,
   const std::int64_t size_global = bs * map.size_global();
   const std::vector<PetscInt> ghosts(map.ghosts().begin(), map.ghosts().end());
   Vec vec;
-  VecCreateGhostBlockWithArray(map.comm(), bs, size_local, size_global,
-                               ghosts.size(), ghosts.data(), x.data(), &vec);
+  PetscErrorCode ierr;
+  if (bs == 1)
+  {
+    ierr
+        = VecCreateGhostWithArray(map.comm(), size_local, size_global,
+                                  ghosts.size(), ghosts.data(), x.data(), &vec);
+    CHECK_ERROR("VecCreateGhostWithArray");
+  }
+  else
+  {
+    ierr = VecCreateGhostBlockWithArray(map.comm(), bs, size_local, size_global,
+                                        ghosts.size(), ghosts.data(), x.data(),
+                                        &vec);
+    CHECK_ERROR("VecCreateGhostBlockWithArray");
+  }
+
+  assert(vec);
   return vec;
 }
 //-----------------------------------------------------------------------------
@@ -368,7 +392,6 @@ void petsc::options::clear()
 }
 //-----------------------------------------------------------------------------
 //-----------------------------------------------------------------------------
-//-----------------------------------------------------------------------------
 petsc::Vector::Vector(const common::IndexMap& map, int bs)
     : _x(la::petsc::create_vector(map, bs))
 {
@@ -688,18 +711,6 @@ int petsc::KrylovSolver::solve(Vec x, const Vec b, bool transpose) const
       petsc::error(ierr, __FILE__, "KSPSolve");
   }
 
-  // FIXME: Remove ghost updating?
-  // Update ghost values in solution vector
-  Vec xg;
-  VecGhostGetLocalForm(x, &xg);
-  const bool is_ghosted = xg ? true : false;
-  VecGhostRestoreLocalForm(x, &xg);
-  if (is_ghosted)
-  {
-    VecGhostUpdateBegin(x, INSERT_VALUES, SCATTER_FORWARD);
-    VecGhostUpdateEnd(x, INSERT_VALUES, SCATTER_FORWARD);
-  }
-
   // Get the number of iterations
   PetscInt num_iterations = 0;
   ierr = KSPGetIterationNumber(_ksp, &num_iterations);

diff --git a/python/demo/demo_axis/demo_axis.py b/python/demo/demo_axis/demo_axis.py
@@ -18,15 +18,14 @@
 from functools import partial
 
 import numpy as np
-from mesh_sphere_axis import generate_mesh_sphere_axis
-from scipy.special import jv, jvp
-
 import ufl
 from basix.ufl import element, mixed_element
-from dolfinx import fem, io, mesh, plot
-
+from dolfinx.fem.petsc import LinearProblem
+from mesh_sphere_axis import generate_mesh_sphere_axis
 from mpi4py import MPI
-from petsc4py import PETSc
+from scipy.special import jv, jvp
+
+from dolfinx import default_scalar_type, fem, io, mesh, plot
 
 try:
     from dolfinx.io import VTXWriter
@@ -52,7 +51,7 @@
 # The time-harmonic Maxwell equation is complex-valued PDE. PETSc must
 # therefore have compiled with complex scalars.
 
-if not np.issubdtype(PETSc.ScalarType, np.complexfloating):
+if not np.issubdtype(default_scalar_type, np.complexfloating):
     print("Demo should only be executed with DOLFINx complex mode")
     exit(0)
 
@@ -491,7 +490,7 @@ def create_eps_mu(pml, rho, eps_bkg, mu_bkg):
 phi = np.pi / 4
 
 # Initialize phase term
-phase = fem.Constant(msh, PETSc.ScalarType(np.exp(1j * 0 * phi)))
+phase = fem.Constant(msh, default_scalar_type(np.exp(1j * 0 * phi)))
 # -
 
 # We now solve the problem:
@@ -518,7 +517,7 @@ def create_eps_mu(pml, rho, eps_bkg, mu_bkg):
         + k0 ** 2 * ufl.inner(eps_pml * Es_m, v_m) * rho * dPml
     a, L = ufl.lhs(F), ufl.rhs(F)
 
-    problem = fem.petsc.LinearProblem(a, L, bcs=[], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
+    problem = LinearProblem(a, L, bcs=[], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
     Esh_m = problem.solve()
 
     # Scattered magnetic field

diff --git a/python/demo/demo_biharmonic.py b/python/demo/demo_biharmonic.py
@@ -115,6 +115,7 @@
 
 import ufl
 from dolfinx import fem, io, mesh, plot
+from dolfinx.fem.petsc import LinearProblem
 from dolfinx.mesh import CellType, GhostMode
 from ufl import (CellDiameter, FacetNormal, avg, div, dS, dx, grad, inner,
                  jump, pi, sin)
@@ -208,14 +209,13 @@
 L = inner(f, v) * dx
 # -
 
-# We create a {py:class}`LinearProblem <dolfinx.fem.LinearProblem>`
+# We create a {py:class}`LinearProblem <dolfinx.fem.petsc.LinearProblem>`
 # object that brings together the variational problem, the Dirichlet
 # boundary condition, and which specifies the linear solver. In this
 # case we use a direct (LU) solver. The {py:func}`solve
-# <dolfinx.fem.LinearProblem.solve>` will compute a solution.
+# <dolfinx.fem.petsc.LinearProblem.solve>` will compute a solution.
 
-problem = fem.petsc.LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly",
-                                                                 "pc_type": "lu"})
+problem = LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
 uh = problem.solve()
 
 # The solution can be written to a  {py:class}`XDMFFile

diff --git a/python/demo/demo_elasticity.py b/python/demo/demo_elasticity.py
@@ -33,11 +33,12 @@
 from dolfinx.io import XDMFFile
 from dolfinx.mesh import (CellType, GhostMode, create_box,
                           locate_entities_boundary)
+from mpi4py import MPI
 from petsc4py import PETSc
 from ufl import dx, grad, inner
 
-from dolfinx import default_real_type, la
-from mpi4py import MPI
+import dolfinx
+from dolfinx import la
 
 dtype = PETSc.ScalarType
 # -
@@ -58,33 +59,32 @@ def build_nullspace(V):
     # Create vectors that will span the nullspace
     bs = V.dofmap.index_map_bs
     length0 = V.dofmap.index_map.size_local
-    length1 = length0 + V.dofmap.index_map.num_ghosts
-    basis = [np.zeros(bs * length1, dtype=dtype) for i in range(6)]
+    basis = [la.vector(V.dofmap.index_map, bs=bs, dtype=dtype) for i in range(6)]
+    b = [b.array for b in basis]
 
     # Get dof indices for each subspace (x, y and z dofs)
     dofs = [V.sub(i).dofmap.list.flatten() for i in range(3)]
 
     # Set the three translational rigid body modes
     for i in range(3):
-        basis[i][dofs[i]] = 1.0
+        b[i][dofs[i]] = 1.0
 
     # Set the three rotational rigid body modes
     x = V.tabulate_dof_coordinates()
     dofs_block = V.dofmap.list.flatten()
     x0, x1, x2 = x[dofs_block, 0], x[dofs_block, 1], x[dofs_block, 2]
-    basis[3][dofs[0]] = -x1
-    basis[3][dofs[1]] = x0
-    basis[4][dofs[0]] = x2
-    basis[4][dofs[2]] = -x0
-    basis[5][dofs[2]] = x1
-    basis[5][dofs[1]] = -x2
-
-    # Create PETSc Vec objects (excluding ghosts) and normalise
-    basis_petsc = [PETSc.Vec().createWithArray(x[:bs * length0], bsize=3, comm=V.mesh.comm) for x in basis]
-    la.orthonormalize(basis_petsc)
-    assert la.is_orthonormal(basis_petsc, eps=1.0e3 * np.finfo(default_real_type).eps)
-
-    # Create and return a PETSc nullspace
+    b[3][dofs[0]] = -x1
+    b[3][dofs[1]] = x0
+    b[4][dofs[0]] = x2
+    b[4][dofs[2]] = -x0
+    b[5][dofs[2]] = x1
+    b[5][dofs[1]] = -x2
+
+    _basis = [x._cpp_object for x in basis]
+    dolfinx.cpp.la.orthonormalize(_basis)
+    assert dolfinx.cpp.la.is_orthonormal(_basis)
+
+    basis_petsc = [PETSc.Vec().createWithArray(x[:bs * length0], bsize=3, comm=V.mesh.comm) for x in b]
     return PETSc.NullSpace().create(vectors=basis_petsc)
 
 
@@ -167,6 +167,7 @@ def σ(v):
 
 ns = build_nullspace(V)
 A.setNearNullSpace(ns)
+A.setOption(PETSc.Mat.Option.SPD, True)
 
 # Set PETSc solver options, create a PETSc Krylov solver, and attach the
 # matrix `A` to the solver:
@@ -175,16 +176,15 @@ def σ(v):
 # Set solver options
 opts = PETSc.Options()
 opts["ksp_type"] = "cg"
-opts["ksp_rtol"] = 1.0e-10
+opts["ksp_rtol"] = 1.0e-8
 opts["pc_type"] = "gamg"
 
 # Use Chebyshev smoothing for multigrid
 opts["mg_levels_ksp_type"] = "chebyshev"
 opts["mg_levels_pc_type"] = "jacobi"
 
 # Improve estimate of eigenvalues for Chebyshev smoothing
-opts["mg_levels_esteig_ksp_type"] = "cg"
-opts["mg_levels_ksp_chebyshev_esteig_steps"] = 20
+opts["mg_levels_ksp_chebyshev_esteig_steps"] = 10
 
 # Create PETSc Krylov solver and turn convergence monitoring on
 solver = PETSc.KSP().create(msh.comm)

diff --git a/python/demo/demo_lagrange_variants.py b/python/demo/demo_lagrange_variants.py
@@ -17,21 +17,17 @@
 #
 # We begin this demo by importing the required modules.
 
-import matplotlib.pylab as plt
-import numpy as np
-
 # +
 import basix
 import basix.ufl
+import matplotlib.pylab as plt
+import numpy as np
 import ufl
-from dolfinx import default_scalar_type, fem, mesh
-from ufl import ds, dx, grad, inner
-
+from dolfinx.fem.petsc import LinearProblem
 from mpi4py import MPI
+from ufl import ds, dx, grad, inner
 
-if np.issubdtype(default_scalar_type, np.complexfloating):
-    print("Demo should only be executed with DOLFINx real mode")
-    exit(0)
+from dolfinx import default_scalar_type, fem, mesh
 # -
 
 # Note that Basix and the Basix UFL wrapper are imported directly. Basix
@@ -132,7 +128,7 @@
 a = inner(grad(u), grad(v)) * dx
 L = inner(f, v) * dx + inner(g, v) * ds
 
-problem = fem.petsc.LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
+problem = LinearProblem(a, L, bcs=[bc], petsc_options={"ksp_type": "preonly", "pc_type": "lu"})
 uh = problem.solve()
 # -