JDBetteridge/petsc priorities (#3348)

* Move all PETSc defaults to petsc.py

* Add warnings when defaulting to PETSc internals

* Stop installing Chaco by default

* Remove Chaco and ML from docker container

---------

Co-authored-by: David A. Ham <david.ham@imperial.ac.uk>
Co-authored-by: Josh Hope-Collins <joshua.hope-collins13@imperial.ac.uk>

docker/Dockerfile.env

@@ -47,13 +47,11 @@ RUN bash -c 'cd petsc; \
         --with-make-np=12 \
         --with-shared-libraries=1 \
         --with-zlib \
-        --download-chaco \
         --download-fftw \
         --download-hdf5 \
         --download-hwloc \
         --download-hypre \
         --download-metis \
-        --download-ml \
         --download-mumps \
         --download-mpich \
         --download-mpich-device=ch3:sock \
@@ -84,13 +82,11 @@ RUN bash -c 'export PACKAGES=/home/firedrake/petsc/packages; \
         --with-bison \
         --with-flex \
         --with-zlib \
-        --with-chaco-dir=$PACKAGES \
         --with-fftw-dir=$PACKAGES \
         --with-hdf5-dir=$PACKAGES \
         --with-hwloc-dir=$PACKAGES \
         --with-hypre-dir=$PACKAGES \
         --with-metis-dir=$PACKAGES \
-        --with-ml-dir=$PACKAGES \
         --with-mpi-dir=$PACKAGES \
         --with-mumps-dir=$PACKAGES \
         --with-netcdf-dir=$PACKAGES \
@@ -126,7 +122,6 @@ RUN bash -c 'export PACKAGES=/home/firedrake/petsc/packages; \
         --with-bison \
         --with-flex \
         --with-zlib \
-        --with-chaco-dir=$PACKAGES \
         --with-fftw-dir=$PACKAGES \
         --with-hdf5-dir=$PACKAGES \
         --with-hwloc-dir=$PACKAGES \

firedrake/function.py

@@ -642,9 +642,9 @@ def at(self, arg, *args, **kwargs):
             raise ValueError("Point dimension (%d) does not match geometric dimension (%d)." % (arg.shape[-1], gdim))
 
         # Check if we have got the same points on each process
-        root_arg = self.comm.bcast(arg, root=0)
+        root_arg = self._comm.bcast(arg, root=0)
         same_arg = arg.shape == root_arg.shape and np.allclose(arg, root_arg)
-        diff_arg = self.comm.allreduce(int(not same_arg), op=MPI.SUM)
+        diff_arg = self._comm.allreduce(int(not same_arg), op=MPI.SUM)
         if diff_arg:
             raise ValueError("Points to evaluate are inconsistent among processes.")
 

firedrake/mesh.py

@@ -31,7 +31,9 @@
 from firedrake.utils import IntType, RealType
 from firedrake.logging import info_red
 from firedrake.parameters import parameters
-from firedrake.petsc import PETSc, OptionsManager
+from firedrake.petsc import (
+    PETSc, OptionsManager, get_external_packages, DEFAULT_PARTITIONER
+)
 from firedrake.adjoint_utils import MeshGeometryMixin
 from pyadjoint import stop_annotating
 
@@ -1228,38 +1230,35 @@ def _set_partitioner(self, plex, distribute, partitioner_type=None):
             "shell", or `None` (unspecified). Ignored if the distribute parameter
             specifies the distribution.
         """
-        from firedrake_configuration import get_config
         if plex.comm.size == 1 or distribute is False:
             return
         partitioner = plex.getPartitioner()
         if distribute is True:
             if partitioner_type:
-                if partitioner_type not in ["chaco", "ptscotch", "parmetis"]:
-                    raise ValueError("Unexpected partitioner_type %s" % partitioner_type)
-                if partitioner_type == "chaco":
-                    if IntType.itemsize == 8:
-                        raise ValueError("Unable to use 'chaco': 'chaco' is 32 bit only, "
-                                         "but your Integer is %d bit." % IntType.itemsize * 8)
-                    if plex.isDistributed():
-                        raise ValueError("Unable to use 'chaco': 'chaco' is a serial "
-                                         "patitioner, but the mesh is distributed.")
-                if partitioner_type == "parmetis":
-                    if not get_config().get("options", {}).get("with_parmetis", False):
-                        raise ValueError("Unable to use 'parmetis': Firedrake is not "
-                                         "installed with 'parmetis'.")
+                if partitioner_type not in ["chaco", "ptscotch", "parmetis", "simple", "shell"]:
+                    raise ValueError(
+                        f"Unexpected partitioner_type: {partitioner_type}")
+                if partitioner_type in ["chaco", "ptscotch", "parmetis"] and \
+                        partitioner_type not in get_external_packages():
+                    raise ValueError(
+                        f"Unable to use {partitioner_type} as PETSc is not "
+                        f"installed with {partitioner_type}."
+                    )
+                if partitioner_type == "chaco" and plex.isDistributed():
+                    raise ValueError(
+                        "Unable to use 'chaco' mesh partitioner, 'chaco' is a "
+                        "serial partitioner, but the mesh is distributed."
+                    )
             else:
-                if IntType.itemsize == 8 or plex.isDistributed():
-                    # Default to PTSCOTCH on 64bit ints (Chaco is 32 bit int only).
-                    # Chaco does not work on distributed meshes.
-                    if get_config().get("options", {}).get("with_parmetis", False):
-                        partitioner_type = "parmetis"
-                    else:
-                        partitioner_type = "ptscotch"
-                else:
-                    partitioner_type = "chaco"
-            partitioner.setType({"chaco": partitioner.Type.CHACO,
-                                 "ptscotch": partitioner.Type.PTSCOTCH,
-                                 "parmetis": partitioner.Type.PARMETIS}[partitioner_type])
+                partitioner_type = DEFAULT_PARTITIONER
+
+            partitioner.setType({
+                "chaco": partitioner.Type.CHACO,
+                "ptscotch": partitioner.Type.PTSCOTCH,
+                "parmetis": partitioner.Type.PARMETIS,
+                "shell": partitioner.Type.SHELL,
+                "simple": partitioner.Type.SIMPLE
+            }[partitioner_type])
         else:
             sizes, points = distribute
             partitioner.setType(partitioner.Type.SHELL)

firedrake/petsc.py

@@ -4,15 +4,24 @@
 import os
 import subprocess
 from contextlib import contextmanager
+from copy import deepcopy
+from types import MappingProxyType
 from typing import Any
+from warnings import warn
 
 import petsc4py
 from mpi4py import MPI
 from petsc4py import PETSc
 from pyop2 import mpi
 
 
-__all__ = ("PETSc", "OptionsManager", "get_petsc_variables")
+__all__ = (
+    "PETSc",
+    "OptionsManager",
+    "get_petsc_variables",
+    "get_petscconf_h",
+    "get_external_packages"
+)
 
 
 class FiredrakePETScError(Exception):
@@ -102,6 +111,37 @@ def get_petsc_variables():
     return {k.strip(): v.strip() for k, v in splitlines}
 
 
+@functools.lru_cache()
+def get_petscconf_h():
+    """Get dict of PETSc include variables from the file:
+    $PETSC_DIR/$PETSC_ARCH/include/petscconf.h
+
+    The ``#define`` and ``PETSC_`` prefix are dropped in the dictionary key.
+
+    The result is memoized to avoid constantly reading the file.
+    """
+    config = petsc4py.get_config()
+    path = [config["PETSC_DIR"], config["PETSC_ARCH"], "include/petscconf.h"]
+    petscconf_h = os.path.join(*path)
+    with open(petscconf_h) as fh:
+        # TODO: use `removeprefix("#define PETSC_")` in place of
+        # `lstrip("#define PETSC")[1:]` when support for Python 3.8 is dropped
+        splitlines = (
+            line.lstrip("#define PETSC")[1:].split(" ", maxsplit=1)
+            for line in filter(lambda x: x.startswith("#define PETSC_"), fh.readlines())
+        )
+    return {k: v.strip() for k, v in splitlines}
+
+
+def get_external_packages():
+    """Return a list of PETSc external packages that are installed.
+
+    """
+    # The HAVE_PACKAGES variable uses delimiters at both ends
+    # so we drop the empty first and last items
+    return get_petscconf_h()["HAVE_PACKAGES"].split(":")[1:-1]
+
+
 def _get_dependencies(filename):
     """Get all the dependencies of a shared object library"""
     # Linux uses `ldd` to look at shared library linkage, MacOS uses `otool`
@@ -334,3 +374,79 @@ def garbage_view(obj: Any):
         PETSc.garbage_view(comm)
     else:
         raise FiredrakePETScError("No comm found, cannot view garbage")
+
+
+external_packages = get_external_packages()
+
+# Setup default partitioner
+# Manually define the priority until
+# https://petsc.org/main/src/dm/partitioner/interface/partitioner.c.html#PetscPartitionerGetDefaultType
+# is added to petsc4py
+partitioner_priority = ["parmetis", "ptscotch", "chaco"]
+for partitioner in partitioner_priority:
+    if partitioner in external_packages:
+        DEFAULT_PARTITIONER = partitioner
+        break
+else:
+    warn(
+        "No external package for " + ", ".join(partitioner_priority)
+        + " found, defaulting to PETSc simple partitioner. This may not be optimal."
+    )
+    DEFAULT_PARTITIONER = "simple"
+
+# Setup default direct solver
+direct_solver_priority = ["mumps", "superlu_dist", "pastix"]
+for solver in direct_solver_priority:
+    if solver in external_packages:
+        DEFAULT_DIRECT_SOLVER = solver
+        _DEFAULT_DIRECT_SOLVER_PARAMETERS = {"mat_solver_type": solver}
+        break
+else:
+    warn(
+        "No external package for " + ", ".join(direct_solver_priority)
+        + " found, defaulting to PETSc LU. This will only work in serial."
+    )
+    DEFAULT_DIRECT_SOLVER = "petsc"
+    _DEFAULT_DIRECT_SOLVER_PARAMETERS = {"mat_solver_type": "petsc"}
+
+# MUMPS needs an additional parameter set
+# From the MUMPS documentation:
+# > ICNTL(14) controls the percentage increase in the estimated working space...
+# > ... Remarks: When significant extra fill-in is caused by numerical pivoting, increasing ICNTL(14) may help.
+if DEFAULT_DIRECT_SOLVER == "mumps":
+    _DEFAULT_DIRECT_SOLVER_PARAMETERS["mat_mumps_icntl_14"] = 200
+
+# Setup default AMG preconditioner
+amg_priority = ["hypre", "ml"]
+for amg in amg_priority:
+    if amg in external_packages:
+        DEFAULT_AMG_PC = amg
+        break
+else:
+    DEFAULT_AMG_PC = "gamg"
+
+
+# Parameters must be flattened for `set_defaults` in `solving_utils.py` to
+# mutate options dictionaries "correctly".
+# TODO: refactor `set_defaults` in `solving_utils.py`
+_DEFAULT_KSP_PARAMETERS = flatten_parameters({
+    "mat_type": "aij",
+    "ksp_type": "preonly",
+    "ksp_rtol": 1e-7,
+    "pc_type": "lu",
+    "pc_factor": _DEFAULT_DIRECT_SOLVER_PARAMETERS
+})
+
+_DEFAULT_SNES_PARAMETERS = {
+    "snes_type": "newtonls",
+    "snes_linesearch_type": "basic",
+    # Really we want **DEFAULT_KSP_PARAMETERS in here, but it isn't the way the NonlinearVariationalSovler class works
+}
+# We also want looser KSP tolerances for non-linear solves
+# DEFAULT_SNES_PARAMETERS["ksp_rtol"] = 1e-5
+# this is specified in the NonlinearVariationalSolver class
+
+# Make all of the `DEFAULT_` dictionaries immutable so someone doesn't accidentally overwrite them
+DEFAULT_DIRECT_SOLVER_PARAMETERS = MappingProxyType(deepcopy(_DEFAULT_DIRECT_SOLVER_PARAMETERS))
+DEFAULT_KSP_PARAMETERS = MappingProxyType(deepcopy(_DEFAULT_KSP_PARAMETERS))
+DEFAULT_SNES_PARAMETERS = MappingProxyType(deepcopy(_DEFAULT_SNES_PARAMETERS))

firedrake/solving_utils.py

@@ -3,10 +3,9 @@
 import numpy
 
 from pyop2 import op2
-from firedrake_configuration import get_config
 from firedrake import function, cofunction, dmhooks
 from firedrake.exceptions import ConvergenceError
-from firedrake.petsc import PETSc
+from firedrake.petsc import PETSc, DEFAULT_KSP_PARAMETERS
 from firedrake.formmanipulation import ExtractSubBlock
 from firedrake.utils import cached_property
 from firedrake.logging import warning
@@ -18,33 +17,20 @@ def _make_reasons(reasons):
 
 
 KSPReasons = _make_reasons(PETSc.KSP.ConvergedReason())
-
-
 SNESReasons = _make_reasons(PETSc.SNES.ConvergedReason())
 
 
-if get_config()["options"]["petsc_int_type"] == "int32":
-    DEFAULT_KSP_PARAMETERS = {"mat_type": "aij",
-                              "ksp_type": "preonly",
-                              "ksp_rtol": 1e-7,
-                              "pc_type": "lu",
-                              "pc_factor_mat_solver_type": "mumps",
-                              "mat_mumps_icntl_14": 200}
-else:
-    DEFAULT_KSP_PARAMETERS = {"mat_type": "aij",
-                              "ksp_type": "preonly",
-                              "ksp_rtol": 1e-7,
-                              "pc_type": "lu",
-                              "pc_factor_mat_solver_type": "superlu_dist"}
-
-
-def set_defaults(solver_parameters, arguments, *, ksp_defaults={}, snes_defaults={}):
+def set_defaults(solver_parameters, arguments, *, ksp_defaults=None, snes_defaults=None):
     """Set defaults for solver parameters.
 
     :arg solver_parameters: dict of user solver parameters to override/extend defaults
     :arg arguments: arguments for the bilinear form (need to know if we have a Real block).
     :arg ksp_defaults: Default KSP parameters.
     :arg snes_defaults: Default SNES parameters."""
+    if ksp_defaults is None:
+        ksp_defaults = {}
+    if snes_defaults is None:
+        snes_defaults = {}
     if solver_parameters:
         # User configured something, try and set sensible direct solve
         # defaults for missing bits.

firedrake/variational_solver.py

@@ -1,10 +1,14 @@
 import ufl
 from itertools import chain
 from contextlib import ExitStack
+from types import MappingProxyType
 
 from firedrake import dmhooks, slate, solving, solving_utils, ufl_expr, utils
 from firedrake import function
-from firedrake.petsc import PETSc, OptionsManager, flatten_parameters
+from firedrake.petsc import (
+    PETSc, OptionsManager, flatten_parameters, DEFAULT_KSP_PARAMETERS,
+    DEFAULT_SNES_PARAMETERS
+)
 from firedrake.function import Function
 from firedrake.functionspace import RestrictedFunctionSpace
 from firedrake.ufl_expr import TrialFunction, TestFunction
@@ -126,12 +130,10 @@ def dm(self):
 class NonlinearVariationalSolver(OptionsManager, NonlinearVariationalSolverMixin):
     r"""Solves a :class:`NonlinearVariationalProblem`."""
 
-    DEFAULT_SNES_PARAMETERS = {"snes_type": "newtonls",
-                               "snes_linesearch_type": "basic"}
+    DEFAULT_SNES_PARAMETERS = DEFAULT_SNES_PARAMETERS
 
     # Looser default tolerance for KSP inside SNES.
-    DEFAULT_KSP_PARAMETERS = solving_utils.DEFAULT_KSP_PARAMETERS.copy()
-    DEFAULT_KSP_PARAMETERS["ksp_rtol"] = 1e-5
+    DEFAULT_KSP_PARAMETERS = MappingProxyType(DEFAULT_KSP_PARAMETERS | {'ksp_rtol': 1e-5})
 
     @PETSc.Log.EventDecorator()
     @NonlinearVariationalSolverMixin._ad_annotate_init
@@ -402,7 +404,8 @@ class LinearVariationalSolver(NonlinearVariationalSolver):
 
     DEFAULT_SNES_PARAMETERS = {"snes_type": "ksponly"}
 
-    DEFAULT_KSP_PARAMETERS = solving_utils.DEFAULT_KSP_PARAMETERS
+    # Tighter default tolerance for KSP only.
+    DEFAULT_KSP_PARAMETERS = DEFAULT_KSP_PARAMETERS
 
     def invalidate_jacobian(self):
         r"""

scripts/firedrake-install

@@ -707,19 +707,11 @@ def get_minimal_petsc_packages():
     pkgs = set()
     # File format
     pkgs.add("hdf5")
-    # Sparse direct solver
-    pkgs.add("superlu_dist")
     # Parallel mesh partitioner
     pkgs.add("ptscotch")
     # Sparse direct solver
     pkgs.add("scalapack")  # Needed for mumps
     pkgs.add("mumps")
-    if not options["complex"]:
-        # AMG
-        pkgs.add("hypre")
-    if options["petsc_int_type"] == "int32":
-        # Serial mesh partitioner
-        pkgs.add("chaco")
     return pkgs
 
 
@@ -744,6 +736,8 @@ def get_petsc_options(minimal=False):
         petsc_options.add("--download-cmake")
 
     if (not options["minimal_petsc"]) and (not minimal):
+        # Another sparse direct solver
+        petsc_options.add("--download-superlu_dist")
         petsc_options.add("--with-zlib")
         # File formats
         petsc_options.add("--download-netcdf")
@@ -766,6 +760,9 @@ def get_petsc_options(minimal=False):
 
     if options["complex"]:
         petsc_options.add('--with-scalar-type=complex')
+    else:
+        # AMG
+        petsc_options.add("hypre")
 
     if options.get("mpiexec") is not None:
         petsc_options.add("--with-mpiexec={}".format(options["mpiexec"]))