Source code for firedrake.petsc

import functools
import gc
import itertools
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__ = (

class FiredrakePETScError(Exception):

def flatten_parameters(parameters, sep="_"):
    """Flatten a nested parameters dict, joining keys with sep.

    :arg parameters: a dict to flatten.
    :arg sep: separator of keys.

    Used to flatten parameter dictionaries with nested structure to a
    flat dict suitable to pass to PETSc.  For example:

    .. code-block:: python3

       flatten_parameters({"a": {"b": {"c": 4}, "d": 2}, "e": 1}, sep="_")
       => {"a_b_c": 4, "a_d": 2, "e": 1}

    If a "prefix" key already ends with the provided separator, then
    it is not used to concatenate the keys.  Hence:

    .. code-block:: python3

       flatten_parameters({"a_": {"b": {"c": 4}, "d": 2}, "e": 1}, sep="_")
       => {"a_b_c": 4, "a_d": 2, "e": 1}
       # rather than
       => {"a__b_c": 4, "a__d": 2, "e": 1}
    from firedrake.logging import warning
    new = type(parameters)()

    if not len(parameters):
        return new

    def flatten(parameters, *prefixes):
        """Iterate over nested dicts, yielding (*keys, value) pairs."""
        sentinel = object()
            option = sentinel
            for option, value in parameters.items():
                # Recurse into values to flatten any dicts.
                for pair in flatten(value, option, *prefixes):
                    yield pair
            # Make sure zero-length dicts come back.
            if option is sentinel:
                yield (prefixes, parameters)
        except AttributeError:
            # Non dict values are just returned.
            yield (prefixes, parameters)

    def munge(keys):
        """Ensure that each intermediate key in keys ends in sep.

        Also, reverse the list."""
        for key in reversed(keys[1:]):
            if len(key) and not key.endswith(sep):
                yield key + sep
                yield key
            yield keys[0]

    for keys, value in flatten(parameters):
        option = "".join(map(str, munge(keys)))
        if option in new:
            warning("Ignoring duplicate option: %s (existing value %s, new value %s)",
                    option, new[option], value)
        new[option] = value
    return new

[docs] @functools.lru_cache() def get_petsc_variables(): """Get dict of PETSc environment variables from the file: $PETSC_DIR/$PETSC_ARCH/lib/petsc/conf/petscvariables The result is memoized to avoid constantly reading the file. """ config = petsc4py.get_config() path = [config["PETSC_DIR"], config["PETSC_ARCH"], "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}
[docs] @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}
[docs] 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` try: program = ["ldd"] cmd =[*program, filename], stdout=subprocess.PIPE) # Filter out the VDSO and the ELF interpreter on Linux results = [line for line in cmd.stdout.decode("utf-8").split("\n") if "=>" in line] return [line.split()[2] for line in results] except FileNotFoundError: program = ["otool", "-L"] cmd =[*program, filename], stdout=subprocess.PIPE) # Meanwhile MacOS puts garbage at the beginning and end of `otool` output return [line.split()[0] for line in cmd.stdout.decode("utf-8").split("\n")[1:-1]] def get_blas_library(): """Get the path to the BLAS library that PETSc links to""" petsc_py_dependencies = _get_dependencies(PETSc.__file__) library_names = ["blas", "libmkl"] for filename in petsc_py_dependencies: if any(name in filename for name in library_names): return filename # On newer MacOS versions, the PETSc Python extension library doesn't link # to BLAS or MKL directly, so we check the PETSc C library. petsc_c_library = [f for f in petsc_py_dependencies if "libpetsc" in f][0] petsc_c_dependencies = _get_dependencies(petsc_c_library) for filename in petsc_c_dependencies: if any(name in filename for name in library_names): return filename return None
[docs] class OptionsManager(object): # What appeared on the commandline, we should never clear these. # They will override options passed in as a dict if an # options_prefix was supplied. commandline_options = frozenset(PETSc.Options().getAll()) options_object = PETSc.Options() count = itertools.count() """Mixin class that helps with managing setting petsc options. :arg parameters: The dictionary of parameters to use. :arg options_prefix: The prefix to look up items in the global options database (may be ``None``, in which case only entries from ``parameters`` will be considered. If no trailing underscore is provided, one is appended. Hence ``foo_`` and ``foo`` are treated equivalently. As an exception, if the prefix is the empty string, no underscore is appended. To use this, you must call its constructor to with the parameters you want in the options database. You then call :meth:`set_from_options`, passing the PETSc object you'd like to call ``setFromOptions`` on. Note that this will actually only call ``setFromOptions`` the first time (so really this parameters object is a once-per-PETSc-object thing). So that the runtime monitors which look in the options database actually see options, you need to ensure that the options database is populated at the time of a ``SNESSolve`` or ``KSPSolve`` call. Do that using the :meth:`inserted_options` context manager. .. code-block:: python3 with self.inserted_options(): self.snes.solve(...) This ensures that the options database has the relevant entries for the duration of the ``with`` block, before removing them afterwards. This is a much more robust way of dealing with the fixed-size options database than trying to clear it out using destructors. This object can also be used only to manage insertion and deletion into the PETSc options database, by using the context manager. """ def __init__(self, parameters, options_prefix): super().__init__() if parameters is None: parameters = {} else: # Convert nested dicts parameters = flatten_parameters(parameters) if options_prefix is None: self.options_prefix = "firedrake_%d_" % next(self.count) self.parameters = parameters self.to_delete = set(parameters) else: if len(options_prefix) and not options_prefix.endswith("_"): options_prefix += "_" self.options_prefix = options_prefix # Remove those options from the dict that were passed on # the commandline. self.parameters = {k: v for k, v in parameters.items() if options_prefix + k not in self.commandline_options} self.to_delete = set(self.parameters) # Now update parameters from options, so that they're # available to solver setup (for, e.g., matrix-free). # Can't ask for the prefixed guy in the options object, # since that does not DTRT for flag options. for k, v in self.options_object.getAll().items(): if k.startswith(self.options_prefix): self.parameters[k[len(self.options_prefix):]] = v self._setfromoptions = False
[docs] def set_default_parameter(self, key, val): """Set a default parameter value. :arg key: The parameter name :arg val: The parameter value. Ensures that the right thing happens cleaning up the options database. """ k = self.options_prefix + key if k not in self.options_object and key not in self.parameters: self.parameters[key] = val self.to_delete.add(key)
[docs] def set_from_options(self, petsc_obj): """Set up petsc_obj from the options database. :arg petsc_obj: The PETSc object to call setFromOptions on. Matt says: "Only ever call setFromOptions once". This function ensures we do so. """ if not self._setfromoptions: with self.inserted_options(): petsc_obj.setOptionsPrefix(self.options_prefix) # Call setfromoptions inserting appropriate options into # the options database. petsc_obj.setFromOptions() self._setfromoptions = True
[docs] @contextmanager def inserted_options(self): """Context manager inside which the petsc options database contains the parameters from this object.""" try: for k, v in self.parameters.items(): self.options_object[self.options_prefix + k] = v yield finally: for k in self.to_delete: del self.options_object[self.options_prefix + k]
def _extract_comm(obj: Any) -> MPI.Comm: """Extract and return the Firedrake/PyOP2 internal comm of a given object. Parameters ---------- obj: Any Firedrake object or any comm Returns ------- MPI.Comm Internal communicator """ comm = None # If the object is a communicator check whether it is already an internal # communicator, otherwise get the internal communicator attribute from the # given communicator. if isinstance(obj, (PETSc.Comm, mpi.MPI.Comm)): try: if mpi.is_pyop2_comm(obj): comm = obj else: internal_comm = obj.Get_attr(mpi.innercomm_keyval) if internal_comm is None: comm = obj else: comm = internal_comm except mpi.PyOP2CommError: pass elif hasattr(obj, "_comm"): comm = obj._comm elif hasattr(obj, "comm"): comm = obj.comm return comm @mpi.collective def garbage_cleanup(obj: Any): """Clean up garbage PETSc objects on a Firedrake object or any comm. Parameters ---------- obj: Any Firedrake object with a comm, or any comm """ # We are manually calling the Python cyclic garbage collection routine to # get as many unreachable reference cycles swept up before we call the PETSc # cleanup routine. This routine is designed to free up as much memory as # possible for memory constrained systems gc.collect() comm = _extract_comm(obj) if comm: PETSc.garbage_cleanup(comm) else: raise FiredrakePETScError("No comm found, cannot clean up garbage") @mpi.collective def garbage_view(obj: Any): """View garbage PETSc objects stored on a Firedrake object or any comm. Parameters ---------- obj: Any Firedrake object with a comm, or any comm. """ # We are manually calling the Python cyclic garbage collection routine so # that as many unreachable PETSc objects are visible in the garbage view. gc.collect() comm = _extract_comm(obj) if comm: 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 # # 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 `` to # mutate options dictionaries "correctly". # TODO: refactor `set_defaults` in `` _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))