Use petsctools package

firedrake/__init__.py

@@ -1,20 +1,34 @@
-import sys
-from firedrake.configuration import setup_cache_dirs
+# IMPORTANT: If updating this constraint then corresponding changes may be
+# needed in pyproject.toml and firedrake-configure. Note that unlike in
+# those locations where we want to be strict about the specific version, here
+# we are more permissive. This is to catch the case where users don't update
+# their PETSc for a really long time or accidentally install a too-new release
+# that isn't yet supported.
+PETSC_SUPPORTED_VERSIONS = "<3.23"
+
+
+def init_petsc():
+    import os
+    import sys
+    import petsctools
+
+    # We conditionally pass '-options_left no' as in some circumstances (e.g.
+    # when running pytest) PETSc complains that command line options are not
+    # PETSc options.
+    if os.getenv("FIREDRAKE_DISABLE_OPTIONS_LEFT") == "1":
+        petsctools.init(sys.argv + ["-options_left", "no"], version_spec=PETSC_SUPPORTED_VERSIONS)
+    else:
+        petsctools.init(sys.argv, version_spec=PETSC_SUPPORTED_VERSIONS)
+
+
+# Ensure petsc is initialised right away
+init_petsc()
 
 # Set up the cache directories before importing PyOP2.
+from firedrake.configuration import setup_cache_dirs
+
 setup_cache_dirs()
 
-# Ensure petsc is initialised by us before anything else gets in there.
-# We conditionally pass '-options_left no' as in some circumstances (e.g.
-# when running pytest) PETSc complains that command line options are not
-# PETSc options.
-import os
-import petsc4py
-if os.getenv("FIREDRAKE_DISABLE_OPTIONS_LEFT") == "1":
-    petsc4py.init(sys.argv + ["-options_left", "no"])
-else:
-    petsc4py.init(sys.argv)
-del os, petsc4py
 
 # Initialise PETSc events for both import and entire duration of program
 from firedrake import petsc
@@ -86,6 +100,7 @@
 #   from firedrake._deprecation import output
 #   sys.modules["firedrake.output"] = output
 from firedrake.output import *
+import sys
 sys.modules["firedrake.plot"] = plot
 from firedrake.plot import *
 

pyproject.toml

@@ -20,11 +20,13 @@ dependencies = [
   "mpi4py; python_version < '3.13'",
   "h5py>3.12.1",
   "libsupermesh",
-  # NOTE: If changing the PETSc/SLEPc version then firedrake-configure also needs
-  # changing (as well as other references to petsc4py and slepc4py here)
-  "petsc4py==3.22.2",
   "numpy",
   "packaging",
+  "petsctools @ git+https://github.com/connorjward/petsctools.git",
+  # IMPORTANT: If changing the PETSc/SLEPc version then firedrake-configure and
+  # firedrake/__init__.py also need changing (as well as other references
+  # to petsc4py and slepc4py here).
+  "petsc4py==3.22.2",
   "pkgconfig",
   "progress",
   "pycparser",
@@ -132,9 +134,10 @@ requires = [
   "mpi4py; python_version < '3.13'",
   "numpy",
   "pkgconfig",
+  "petsc4py==3.22.2",
+  "petsctools @ git+https://github.com/connorjward/petsctools.git",
   "pybind11",
   "setuptools>61.2",
-  "petsc4py==3.22.2",
   "rtree>=1.2",
 ]
 build-backend = "setuptools.build_meta"

scripts/firedrake-configure

@@ -39,8 +39,8 @@ ARCH_DEFAULT = FiredrakeArch.DEFAULT
 ARCH_COMPLEX = FiredrakeArch.COMPLEX
 
 
-# NOTE: When updating this variable corresponding changes must be made inside
-# pyproject.toml
+# IMPORTANT: When updating this variable, corresponding changes must be made inside
+# pyproject.toml and firedrake/__init__.py.
 SUPPORTED_PETSC_VERSION = "v3.22.2"
 
 

setup.py

@@ -1,45 +1,21 @@
 import logging
 import os
+import pkgconfig
 import platform
-import sys
-import site
 import shutil
+import site
+from dataclasses import dataclass, field
+from glob import glob
+from pathlib import Path
+
+import libsupermesh
 import numpy as np
 import pybind11
 import petsc4py
+import petsctools
 import rtree
-import libsupermesh
-import pkgconfig
-from dataclasses import dataclass, field
-from setuptools import setup, find_packages, Extension
-from glob import glob
-from pathlib import Path
 from Cython.Build import cythonize
-
-
-petsc_config = petsc4py.get_config()
-
-
-def get_petsc_dir():
-    """Attempts to find the PETSc directory on the system
-    """
-    petsc_dir = petsc_config["PETSC_DIR"]
-    petsc_arch = petsc_config["PETSC_ARCH"]
-    pathlist = [petsc_dir]
-    if petsc_arch:
-        pathlist.append(os.path.join(petsc_dir, petsc_arch))
-    return pathlist
-
-
-def get_petsc_variables():
-    """Attempts obtain a dictionary of PETSc configuration settings
-    """
-    path = [get_petsc_dir()[-1], "lib/petsc/conf/petscvariables"]
-    variables_path = os.path.join(*path)
-    with open(variables_path) as fh:
-        # Split lines on first '=' (assignment)
-        splitlines = (line.split("=", maxsplit=1) for line in fh.readlines())
-    return {k.strip(): v.strip() for k, v in splitlines}
+from setuptools import setup, find_packages, Extension
 
 
 @dataclass
@@ -92,16 +68,18 @@ def __getitem__(self, key):
 # example:
 # gcc -I$PETSC_DIR/include -I$PETSC_DIR/$PETSC_ARCH/include -I/petsc4py/include
 # gcc -L$PETSC_DIR/$PETSC_ARCH/lib -lpetsc -Wl,-rpath,$PETSC_DIR/$PETSC_ARCH/lib
-petsc_dirs = get_petsc_dir()
+petsc_dir = petsctools.get_petsc_dir()
+petsc_arch = petsctools.get_petsc_arch()
+petsc_dirs = [petsc_dir, os.path.join(petsc_dir, petsc_arch)]
 petsc_ = ExternalDependency(
     libraries=["petsc"],
     include_dirs=[petsc4py.get_include()] + [os.path.join(d, "include") for d in petsc_dirs],
     library_dirs=[os.path.join(petsc_dirs[-1], "lib")],
     runtime_library_dirs=[os.path.join(petsc_dirs[-1], "lib")],
 )
-petsc_variables = get_petsc_variables()
-petsc_hdf5_compile_args = petsc_variables.get("HDF5_INCLUDE", "")
-petsc_hdf5_link_args = petsc_variables.get("HDF5_LIB", "")
+petscvariables = petsctools.get_petscvariables()
+petsc_hdf5_compile_args = petscvariables.get("HDF5_INCLUDE", "")
+petsc_hdf5_link_args = petscvariables.get("HDF5_LIB", "")
 
 # HDF5
 # example:
@@ -111,15 +89,15 @@ def __getitem__(self, key):
     # We almost always want to be in this first case!!!
     # PETSc variables only contains the compile/link args, not the paths
     hdf5_ = ExternalDependency(
-        extra_compile_args = petsc_hdf5_compile_args.split(),
-        extra_link_args = petsc_hdf5_link_args.split()
+        extra_compile_args=petsc_hdf5_compile_args.split(),
+        extra_link_args=petsc_hdf5_link_args.split()
     )
 elif os.environ.get("HDF5_DIR"):
     hdf5_dir = Path(os.environ.get("HDF5_DIR"))
     hdf5_ = ExternalDependency(
         libraries=["hdf5"],
-        include_dirs = [str(hdf5_dir.joinpath("include"))],
-        library_dirs = [str(hdf5_dir.joinpath("lib"))]
+        include_dirs=[str(hdf5_dir.joinpath("include"))],
+        library_dirs=[str(hdf5_dir.joinpath("lib"))]
     )
 elif pkgconfig.exists("hdf5"):
     hdf5_ = ExternalDependency(**pkgconfig.parse("hdf5"))
@@ -163,9 +141,10 @@ def __getitem__(self, key):
     libraries=["supermesh"],
 )
 
+
 # The following extensions need to be linked accordingly:
 def extensions():
-    ## CYTHON EXTENSIONS
+    # CYTHON EXTENSIONS
     cython_list = []
     # firedrake/cython/dmcommon.pyx: petsc, numpy
     cython_list.append(Extension(
@@ -223,7 +202,7 @@ def extensions():
         sources=[os.path.join("pyop2", "sparsity.pyx")],
         **(petsc_ + numpy_)
     ))
-    ## PYBIND11 EXTENSIONS
+    # PYBIND11 EXTENSIONS
     pybind11_list = []
     # tinyasm/tinyasm.cpp: petsc, pybind11
     pybind11_list.append(Extension(
@@ -284,9 +263,4 @@ def main():
 
 
 make_firedrake_check_package()
-
-
-setup(
-    packages=find_packages(),
-    ext_modules=extensions()
-)
+setup(packages=find_packages(), ext_modules=extensions())