From 4109283bf24afa06a821bde11a993ac7c81f8e75 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 9 Jan 2025 19:24:14 +0000 Subject: [PATCH 01/10] Fix test for mixed-ness in Function.sub, and fix Cofunction.sub to match (#3961) * fix test for mixed-ness in Function.sub, and fix Cofunction.sub to match * allow FunctionSpace.sub to take negative indices * cofunc docstring --- firedrake/cofunction.py | 28 +++++++++++++++---- firedrake/function.py | 13 +++------ firedrake/functionspaceimpl.py | 8 +----- .../regression/test_vfs_component_bcs.py | 2 +- 4 files changed, 28 insertions(+), 23 deletions(-) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index 4878e6da59..ecfb386dd9 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -4,6 +4,7 @@ from ufl.form import BaseForm from pyop2 import op2, mpi from pyadjoint.tape import stop_annotating, annotate_tape, get_working_tape +from finat.ufl import MixedElement import firedrake.assemble import firedrake.functionspaceimpl as functionspaceimpl from firedrake import utils, vector, ufl_expr @@ -119,11 +120,14 @@ def split(self): @utils.cached_property def _components(self): - if self.function_space().value_size == 1: + if self.function_space().rank == 0: return (self, ) else: - return tuple(type(self)(self.function_space().sub(i), val=op2.DatView(self.dat, i)) - for i in range(self.function_space().value_size)) + if self.dof_dset.cdim == 1: + return (type(self)(self.function_space().sub(0), val=self.dat),) + else: + return tuple(type(self)(self.function_space().sub(i), val=op2.DatView(self.dat, j)) + for i, j in enumerate(np.ndindex(self.dof_dset.dim))) @PETSc.Log.EventDecorator() def sub(self, i): @@ -137,9 +141,9 @@ def sub(self, i): :func:`~.VectorFunctionSpace` or :func:`~.TensorFunctionSpace` this returns a proxy object indexing the ith component of the space, suitable for use in boundary condition application.""" - if len(self.function_space()) == 1: - return self._components[i] - return self.subfunctions[i] + mixed = type(self.function_space().ufl_element()) is MixedElement + data = self.subfunctions if mixed else self._components + return data[i] def function_space(self): r"""Return the :class:`.FunctionSpace`, or :class:`.MixedFunctionSpace` @@ -321,6 +325,12 @@ def vector(self): :class:`Cofunction`""" return vector.Vector(self) + @property + def cell_set(self): + r"""The :class:`pyop2.types.set.Set` of cells for the mesh on which this + :class:`Cofunction` is defined.""" + return self.function_space()._mesh.cell_set + @property def node_set(self): r"""A :class:`pyop2.types.set.Set` containing the nodes of this @@ -330,6 +340,12 @@ def node_set(self): """ return self.function_space().node_set + @property + def dof_dset(self): + r"""A :class:`pyop2.types.dataset.DataSet` containing the degrees of freedom of + this :class:`Cofunction`.""" + return self.function_space().dof_dset + def ufl_id(self): return self.uid diff --git a/firedrake/function.py b/firedrake/function.py index da4d264971..97a08be816 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -15,6 +15,7 @@ from pyop2 import op2, mpi from pyop2.exceptions import DataTypeError, DataValueError +from finat.ufl import MixedElement from firedrake.utils import ScalarType, IntType, as_ctypes from firedrake import functionspaceimpl @@ -147,11 +148,8 @@ def sub(self, i): rank-n :class:`~.FunctionSpace`, this returns a proxy object indexing the ith component of the space, suitable for use in boundary condition application.""" - mixed = len(self.function_space()) != 1 + mixed = type(self.function_space().ufl_element()) is MixedElement data = self.subfunctions if mixed else self._components - bound = len(data) - if i < 0 or i >= bound: - raise IndexError(f"Invalid component {i}, not in [0, {bound})") return data[i] @property @@ -352,11 +350,8 @@ def sub(self, i): :func:`~.VectorFunctionSpace` or :func:`~.TensorFunctionSpace` this returns a proxy object indexing the ith component of the space, suitable for use in boundary condition application.""" - mixed = len(self.function_space()) != 1 + mixed = type(self.function_space().ufl_element()) is MixedElement data = self.subfunctions if mixed else self._components - bound = len(data) - if i < 0 or i >= bound: - raise IndexError(f"Invalid component {i}, not in [0, {bound})") return data[i] @PETSc.Log.EventDecorator() @@ -672,7 +667,7 @@ def single_eval(x, buf): value_shape = self.ufl_shape subfunctions = self.subfunctions - mixed = len(subfunctions) != 1 + mixed = type(self.function_space().ufl_element()) is MixedElement # Local evaluation l_result = [] diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index 8fc81244f7..c366c03b66 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -183,11 +183,8 @@ def _components(self): @PETSc.Log.EventDecorator() def sub(self, i): - mixed = len(self) != 1 + mixed = type(self.ufl_element()) is finat.ufl.MixedElement data = self.subfunctions if mixed else self._components - bound = len(data) - if i < 0 or i >= bound: - raise IndexError(f"Invalid component {i}, not in [0, {bound})") return data[i] @utils.cached_property @@ -664,9 +661,6 @@ def _components(self): def sub(self, i): r"""Return a view into the ith component.""" - bound = len(self._components) - if i < 0 or i >= bound: - raise IndexError(f"Invalid component {i}, not in [0, {bound})") return self._components[i] def __mul__(self, other): diff --git a/tests/firedrake/regression/test_vfs_component_bcs.py b/tests/firedrake/regression/test_vfs_component_bcs.py index 76cbd69ea5..6b6a0ea938 100644 --- a/tests/firedrake/regression/test_vfs_component_bcs.py +++ b/tests/firedrake/regression/test_vfs_component_bcs.py @@ -149,7 +149,7 @@ def test_cant_integrate_subscripted_VFS(V): @pytest.mark.parametrize("cmpt", - [-1, 2]) + [-3, 2]) def test_cant_subscript_outside_components(V, cmpt): with pytest.raises(IndexError): return V.sub(cmpt) From 8e1a7484275dc59fccb581ebece78762e12bb039 Mon Sep 17 00:00:00 2001 From: Daiane Iglesia Dolci <63597005+Ig-dolci@users.noreply.github.com> Date: Fri, 10 Jan 2025 13:06:32 +0000 Subject: [PATCH 02/10] Temp fix for firedrake install on Mac for Xcode16 (#3964) --- scripts/firedrake-install | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/scripts/firedrake-install b/scripts/firedrake-install index 17a8d7fdbc..41bdc21b8f 100755 --- a/scripts/firedrake-install +++ b/scripts/firedrake-install @@ -792,6 +792,10 @@ def get_petsc_options(minimal=False): # on macOS with parallel solvers. Its necessity is unclear, so a # review in the future may be required. petsc_options.add("--LDFLAGS=-Wl,-ld_classic,-dead_strip_dylibs") + if Version(clt["version"]) >= Version("16"): + # Should revisit this flag when the the issue https://gitlab.com/petsc/petsc/-/issues/1692 + # is resolved. + petsc_options.add("--with-cxx-dialect=17") elif options.get("with_blas") is not None: petsc_options.add("--with-blaslapack-dir={}".format(options["with_blas"])) if osname == "Darwin": From a08fc1e6eca598ee4ae3aa97af10537ede3b311f Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Fri, 10 Jan 2025 15:47:38 +0000 Subject: [PATCH 03/10] Untangle compilation process (#3940) --- firedrake/function.py | 5 ++- firedrake/mesh.py | 6 +-- firedrake/mg/kernels.py | 2 +- firedrake/preconditioners/fdm.py | 8 +++- firedrake/preconditioners/patch.py | 13 +++--- firedrake/supermeshing.py | 9 +++-- pyop2/compilation.py | 65 +++++++----------------------- pyop2/global_kernel.py | 24 +++++------ tests/pyop2/test_caching.py | 7 ++-- 9 files changed, 55 insertions(+), 84 deletions(-) diff --git a/firedrake/function.py b/firedrake/function.py index 97a08be816..e1378b21e0 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -770,8 +770,8 @@ def make_c_evaluate(function, c_name="evaluate", ldargs=None, tolerance=None): libspatialindex_so = Path(rtree.core.rt._name).absolute() lsi_runpath = f"-Wl,-rpath,{libspatialindex_so.parent}" ldargs += [str(libspatialindex_so), lsi_runpath] - return compilation.load( - src, "c", c_name, + dll = compilation.load( + src, "c", cppargs=[ f"-I{path.dirname(__file__)}", f"-I{sys.prefix}/include", @@ -780,3 +780,4 @@ def make_c_evaluate(function, c_name="evaluate", ldargs=None, tolerance=None): ldargs=ldargs, comm=function.comm ) + return getattr(dll, c_name) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 9936efcefc..7e9a61773b 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -2688,8 +2688,8 @@ def _c_locator(self, tolerance=None): libspatialindex_so = Path(rtree.core.rt._name).absolute() lsi_runpath = f"-Wl,-rpath,{libspatialindex_so.parent}" - locator = compilation.load( - src, "c", "locator", + dll = compilation.load( + src, "c", cppargs=[ f"-I{os.path.dirname(__file__)}", f"-I{sys.prefix}/include", @@ -2703,7 +2703,7 @@ def _c_locator(self, tolerance=None): ], comm=self.comm ) - + locator = getattr(dll, "locator") locator.argtypes = [ctypes.POINTER(function._CFunction), ctypes.POINTER(ctypes.c_double), ctypes.POINTER(ctypes.c_double), diff --git a/firedrake/mg/kernels.py b/firedrake/mg/kernels.py index f892f6260c..6ce9c4eb55 100644 --- a/firedrake/mg/kernels.py +++ b/firedrake/mg/kernels.py @@ -720,7 +720,7 @@ def name_multiindex(multiindex, name): kernel = lp.make_kernel( domains, instructions, kernel_data, name=kernel_name, target=tsfc.parameters.target, lang_version=(2018, 2)) - kernel = lp.merge([kernel, *subkernels]) + kernel = lp.merge([kernel, *subkernels]).with_entrypoints({kernel_name}) return op2.Kernel( kernel, name=kernel_name, include_dirs=Ainv.include_dirs, headers=Ainv.headers, events=Ainv.events) diff --git a/firedrake/preconditioners/fdm.py b/firedrake/preconditioners/fdm.py index e75172dc8e..51a611ad47 100644 --- a/firedrake/preconditioners/fdm.py +++ b/firedrake/preconditioners/fdm.py @@ -1835,13 +1835,17 @@ def setSubMatCSR(comm, triu=False): return cache.setdefault(key, SparseAssembler.load_setSubMatCSR(comm, triu)) @staticmethod - def load_c_code(code, name, **kwargs): + def load_c_code(code, name, comm, argtypes, restype): petsc_dir = get_petsc_dir() cppargs = [f"-I{d}/include" for d in petsc_dir] ldargs = ([f"-L{d}/lib" for d in petsc_dir] + [f"-Wl,-rpath,{d}/lib" for d in petsc_dir] + ["-lpetsc", "-lm"]) - return load(code, "c", name, cppargs=cppargs, ldargs=ldargs, **kwargs) + dll = load(code, "c", cppargs=cppargs, ldargs=ldargs, comm=comm) + fn = getattr(dll, name) + fn.argtypes = argtypes + fn.restype = restype + return fn @staticmethod def load_setSubMatCSR(comm, triu=False): diff --git a/firedrake/preconditioners/patch.py b/firedrake/preconditioners/patch.py index 0a7bad5575..5e9d0d4fa0 100644 --- a/firedrake/preconditioners/patch.py +++ b/firedrake/preconditioners/patch.py @@ -505,12 +505,13 @@ def load_c_function(code, name, comm): ldargs = (["-L%s/lib" % d for d in get_petsc_dir()] + ["-Wl,-rpath,%s/lib" % d for d in get_petsc_dir()] + ["-lpetsc", "-lm"]) - return load(code, "c", name, - argtypes=[ctypes.c_voidp, ctypes.c_int, ctypes.c_voidp, - ctypes.c_voidp, ctypes.c_voidp, ctypes.c_int, - ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp], - restype=ctypes.c_int, cppargs=cppargs, ldargs=ldargs, - comm=comm) + dll = load(code, "c", cppargs=cppargs, ldargs=ldargs, comm=comm) + fn = getattr(dll, name) + fn.argtypes = [ctypes.c_voidp, ctypes.c_int, ctypes.c_voidp, + ctypes.c_voidp, ctypes.c_voidp, ctypes.c_int, + ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp] + fn.restype = ctypes.c_int + return fn def make_c_arguments(form, kernel, state, get_map, require_state=False, diff --git a/firedrake/supermeshing.py b/firedrake/supermeshing.py index ee576ea6e5..b35cbb0fe4 100644 --- a/firedrake/supermeshing.py +++ b/firedrake/supermeshing.py @@ -434,14 +434,15 @@ def likely(cell_A): includes = ["-I%s/include" % d for d in dirs] libs = ["-L%s/lib" % d for d in dirs] libs = libs + ["-Wl,-rpath,%s/lib" % d for d in dirs] + ["-lpetsc", "-lsupermesh"] - lib = load( - supermesh_kernel_str, "c", "supermesh_kernel", + dll = load( + supermesh_kernel_str, "c", cppargs=includes, ldargs=libs, - argtypes=[ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp], - restype=ctypes.c_int, comm=mesh_A._comm ) + lib = getattr(dll, "supermesh_kernel") + lib.argtypes = [ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp, ctypes.c_voidp] + lib.restype = ctypes.c_int ammm(V_A, V_B, likely, node_locations_A, node_locations_B, M_SS, ctypes.addressof(lib), mat) if orig_value_size == 1: diff --git a/pyop2/compilation.py b/pyop2/compilation.py index d28c945188..a9b724f368 100644 --- a/pyop2/compilation.py +++ b/pyop2/compilation.py @@ -56,7 +56,6 @@ from pyop2.logger import warning, debug, progress, INFO from pyop2.exceptions import CompilationError from pyop2.utils import get_petsc_variables -import pyop2.global_kernel from petsc4py import PETSc @@ -411,51 +410,23 @@ class AnonymousCompiler(Compiler): def load_hashkey(*args, **kwargs): - from pyop2.global_kernel import GlobalKernel - if isinstance(args[0], str): - code_hash = md5(args[0].encode()).hexdigest() - elif isinstance(args[0], GlobalKernel): - code_hash = md5(str(args[0].cache_key).encode()).hexdigest() - else: - pass # This will raise an error in load + code_hash = md5(args[0].encode()).hexdigest() return default_parallel_hashkey(code_hash, *args[1:], **kwargs) @mpi.collective @memory_cache(hashkey=load_hashkey) @PETSc.Log.EventDecorator() -def load(jitmodule, extension, fn_name, cppargs=(), ldargs=(), - argtypes=None, restype=None, comm=None): +def load(code, extension, cppargs=(), ldargs=(), comm=None): """Build a shared library and return a function pointer from it. - :arg jitmodule: The JIT Module which can generate the code to compile, or - the string representing the source code. + :arg code: The code to compile. :arg extension: extension of the source file (c, cpp) - :arg fn_name: The name of the function to return from the resulting library :arg cppargs: A tuple of arguments to the C compiler (optional) :arg ldargs: A tuple of arguments to the linker (optional) - :arg argtypes: A list of ctypes argument types matching the arguments of - the returned function (optional, pass ``None`` for ``void``). This is - only used when string is passed in instead of JITModule. - :arg restype: The return type of the function (optional, pass - ``None`` for ``void``). :kwarg comm: Optional communicator to compile the code on (only rank 0 compiles code) (defaults to pyop2.mpi.COMM_WORLD). """ - if isinstance(jitmodule, str): - class StrCode(object): - def __init__(self, code, argtypes): - self.code_to_compile = code - self.cache_key = (None, code) # We peel off the first - # entry, since for a jitmodule, it's a process-local - # cache key - self.argtypes = argtypes - code = StrCode(jitmodule, argtypes) - elif isinstance(jitmodule, pyop2.global_kernel.GlobalKernel): - code = jitmodule - else: - raise ValueError("Don't know how to compile code of type %r" % type(jitmodule)) - global _compiler if _compiler: # Use the global compiler if it has been set @@ -475,15 +446,7 @@ def __init__(self, code, argtypes): # This call is cached on disk so_name = make_so(compiler_instance, code, extension, comm) # This call might be cached in memory by the OS (system dependent) - dll = ctypes.CDLL(so_name) - - if isinstance(jitmodule, pyop2.global_kernel.GlobalKernel): - _add_profiling_events(dll, code.local_kernel.events) - - fn = getattr(dll, fn_name) - fn.argtypes = code.argtypes - fn.restype = restype - return fn + return ctypes.CDLL(so_name) def expandWl(ldflags): @@ -519,27 +482,27 @@ def setdefault(self, key, default=None): return self[key] -def _make_so_hashkey(compiler, jitmodule, extension, comm): +def _make_so_hashkey(compiler, code, extension, comm): if extension == "cpp": exe = compiler.cxx compiler_flags = compiler.cxxflags else: exe = compiler.cc compiler_flags = compiler.cflags - return (compiler, exe, compiler_flags, compiler.ld, compiler.ldflags, jitmodule.cache_key) + return (compiler, code, exe, compiler_flags, compiler.ld, compiler.ldflags) -def check_source_hashes(compiler, jitmodule, extension, comm): +def check_source_hashes(compiler, code, extension, comm): """A check to see whether code generated on all ranks is identical. :arg compiler: The compiler to use to create the shared library. - :arg jitmodule: The JIT Module which can generate the code to compile. + :arg code: The code to compile. :arg filename: The filename of the library to create. :arg extension: extension of the source file (c, cpp). :arg comm: Communicator over which to perform compilation. """ # Reconstruct hash from filename - hashval = _as_hexdigest(_make_so_hashkey(compiler, jitmodule, extension, comm)) + hashval = _as_hexdigest(_make_so_hashkey(compiler, code, extension, comm)) with mpi.temp_internal_comm(comm) as icomm: matching = icomm.allreduce(hashval, op=_check_op) if matching != hashval: @@ -550,7 +513,7 @@ def check_source_hashes(compiler, jitmodule, extension, comm): output.mkdir(parents=True, exist_ok=True) icomm.barrier() with open(srcfile, "w") as fh: - fh.write(jitmodule.code_to_compile) + fh.write(code) icomm.barrier() raise CompilationError(f"Generated code differs across ranks (see output in {output})") @@ -561,11 +524,11 @@ def check_source_hashes(compiler, jitmodule, extension, comm): cache_factory=lambda: CompilerDiskAccess(configuration['cache_dir'], extension=".so") ) @PETSc.Log.EventDecorator() -def make_so(compiler, jitmodule, extension, comm, filename=None): +def make_so(compiler, code, extension, comm, filename=None): """Build a shared library and load it :arg compiler: The compiler to use to create the shared library. - :arg jitmodule: The JIT Module which can generate the code to compile. + :arg code: The code to compile. :arg filename: The filename of the library to create. :arg extension: extension of the source file (c, cpp). :arg comm: Communicator over which to perform compilation. @@ -605,7 +568,7 @@ def make_so(compiler, jitmodule, extension, comm, filename=None): with progress(INFO, 'Compiling wrapper'): # Write source code to disk with open(cname, "w") as fh: - fh.write(jitmodule.code_to_compile) + fh.write(code) os.close(descriptor) if not compiler.ld: @@ -650,7 +613,7 @@ def _run(cc, logfile, errfile, step="Compilation", filemode="w"): """)) -def _add_profiling_events(dll, events): +def add_profiling_events(dll, events): """ If PyOP2 is in profiling mode, events are attached to dll to profile the local linear algebra calls. The event is generated here in python and then set in the shared library, diff --git a/pyop2/global_kernel.py b/pyop2/global_kernel.py index ae13dc1c59..7edfed0771 100644 --- a/pyop2/global_kernel.py +++ b/pyop2/global_kernel.py @@ -11,7 +11,7 @@ from petsc4py import PETSc from pyop2 import mpi -from pyop2.compilation import load +from pyop2.compilation import add_profiling_events, load from pyop2.configuration import configuration from pyop2.datatypes import IntType, as_ctypes from pyop2.types import IterationRegion, Constant, READ @@ -366,8 +366,11 @@ def code_to_compile(self): """Return the C/C++ source code as a string.""" from pyop2.codegen.rep2loopy import generate - wrapper = generate(self.builder) - code = lp.generate_code_v2(wrapper) + with PETSc.Log.Event("GlobalKernel: generate loopy"): + wrapper = generate(self.builder) + + with PETSc.Log.Event("GlobalKernel: generate device code"): + code = lp.generate_code_v2(wrapper) if self.local_kernel.cpp: from loopy.codegen.result import process_preambles @@ -397,15 +400,12 @@ def compile(self, comm): + tuple(self.local_kernel.ldargs) ) - return load( - self, - extension, - self.name, - cppargs=cppargs, - ldargs=ldargs, - restype=ctypes.c_int, - comm=comm - ) + dll = load(self.code_to_compile, extension, cppargs=cppargs, ldargs=ldargs, comm=comm) + add_profiling_events(dll, self.local_kernel.events) + fn = getattr(dll, self.name) + fn.argtypes = self.argtypes + fn.restype = ctypes.c_int + return fn @cached_property def argtypes(self): diff --git a/tests/pyop2/test_caching.py b/tests/pyop2/test_caching.py index 1298991b3e..cfd9e6ce7f 100644 --- a/tests/pyop2/test_caching.py +++ b/tests/pyop2/test_caching.py @@ -31,7 +31,6 @@ # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED # OF THE POSSIBILITY OF SUCH DAMAGE. -import ctypes import os import pytest import tempfile @@ -785,7 +784,8 @@ def test_writing_large_so(): if COMM_WORLD.rank == 1: os.remove("big.c") - fn = load(program, "c", "big", argtypes=(ctypes.c_voidp,), comm=COMM_WORLD) + dll = load(program, "c", comm=COMM_WORLD) + fn = getattr(dll, "big") assert fn is not None @@ -800,7 +800,8 @@ def test_two_comms_compile_the_same_code(): } """) - fn = load(code, "c", "noop", argtypes=(), comm=COMM_WORLD) + dll = load(code, "c", comm=COMM_WORLD) + fn = getattr(dll, "noop") assert fn is not None From 5c6e320533a9064a55af40e5931655859f5383b4 Mon Sep 17 00:00:00 2001 From: Connor Ward Date: Fri, 10 Jan 2025 17:26:18 +0000 Subject: [PATCH 04/10] Configure PETSc with --with-strict-petscerrorcode (#3962) And necessary fixes to tinyASM --- .github/workflows/pip-mac.yml | 1 + .github/workflows/pyop2.yml | 1 + docker/Dockerfile.env | 3 ++ scripts/firedrake-install | 1 + tinyasm/matinvert.cpp | 2 +- tinyasm/tinyasm.cpp | 78 ++++++++++++++++++----------------- 6 files changed, 47 insertions(+), 39 deletions(-) diff --git a/.github/workflows/pip-mac.yml b/.github/workflows/pip-mac.yml index 54e2ada7cc..cf89b10724 100644 --- a/.github/workflows/pip-mac.yml +++ b/.github/workflows/pip-mac.yml @@ -53,6 +53,7 @@ jobs: --with-shared-libraries=1 \ --with-mpi-dir=/opt/homebrew \ --with-zlib \ + --with-strict-petscerrorcode \ --download-bison \ --download-hdf5 \ --download-hwloc \ diff --git a/.github/workflows/pyop2.yml b/.github/workflows/pyop2.yml index 5104284d05..d6e5a4acfb 100644 --- a/.github/workflows/pyop2.yml +++ b/.github/workflows/pyop2.yml @@ -51,6 +51,7 @@ jobs: --with-debugging=1 \ --with-shared-libraries=1 \ --with-c2html=0 \ + --with-strict-petscerrorcode \ --with-fortran-bindings=0 make diff --git a/docker/Dockerfile.env b/docker/Dockerfile.env index d40a613f85..5a1803a10e 100644 --- a/docker/Dockerfile.env +++ b/docker/Dockerfile.env @@ -66,6 +66,7 @@ RUN bash -c 'cd petsc; \ --download-scalapack \ --download-suitesparse \ --download-superlu_dist \ + --with-strict-petscerrorcode \ PETSC_ARCH=packages; \ mv packages/include/petscconf.h packages/include/old_petscconf.nope; \ rm -rf /home/firedrake/petsc/**/externalpackages; \ @@ -105,6 +106,7 @@ RUN bash -c 'export PACKAGES=/home/firedrake/petsc/packages; \ --with-scalapack-dir=$PACKAGES \ --with-suitesparse-dir=$PACKAGES \ --with-superlu_dist-dir=$PACKAGES \ + --with-strict-petscerrorcode \ PETSC_ARCH=default; \ make PETSC_DIR=/home/firedrake/petsc PETSC_ARCH=default all;' @@ -144,6 +146,7 @@ RUN bash -c 'export PACKAGES=/home/firedrake/petsc/packages; \ --with-scalapack-dir=$PACKAGES \ --with-suitesparse-dir=$PACKAGES \ --with-superlu_dist-dir=$PACKAGES \ + --with-strict-petscerrorcode \ PETSC_ARCH=complex; \ make PETSC_DIR=/home/firedrake/petsc PETSC_ARCH=complex all;' diff --git a/scripts/firedrake-install b/scripts/firedrake-install index 41bdc21b8f..aaf3f293f9 100755 --- a/scripts/firedrake-install +++ b/scripts/firedrake-install @@ -726,6 +726,7 @@ def get_petsc_options(minimal=False): "--with-debugging=0", "--with-shared-libraries=1", "--with-c2html=0", + "--with-strict-petscerrorcode", # Parser generator "--download-bison"} for pkg in get_minimal_petsc_packages(): diff --git a/tinyasm/matinvert.cpp b/tinyasm/matinvert.cpp index 0c1e7025c2..525dcefbd0 100644 --- a/tinyasm/matinvert.cpp +++ b/tinyasm/matinvert.cpp @@ -6,6 +6,6 @@ PetscErrorCode mymatinvert(PetscBLASInt* n, PetscScalar* mat, PetscBLASInt* piv, PetscCheck(!(*info), PETSC_COMM_SELF, PETSC_ERR_LIB, "TinyASM error calling ?getrf in mymatinvert"); PetscCallBLAS("LAPACKgetri", LAPACKgetri_(n, mat, n, piv, work, n, info)); PetscCheck(!(*info), PETSC_COMM_SELF, PETSC_ERR_LIB, "TinyASM error calling ?getri in mymatinvert"); - return PETSC_SUCCESS; + PetscFunctionReturn(PETSC_SUCCESS); } diff --git a/tinyasm/tinyasm.cpp b/tinyasm/tinyasm.cpp index cf583a6fc8..78261b7ac2 100644 --- a/tinyasm/tinyasm.cpp +++ b/tinyasm/tinyasm.cpp @@ -63,31 +63,31 @@ class BlockJacobi { fwork = vector(biggestBlock, 0.); localmats_aij = NULL; dofis = vector(numBlocks); - PetscMalloc1(numBlocks, &localmats); + PetscCallVoid(PetscMalloc1(numBlocks, &localmats)); for(int p=0; p& sf, const std::vector& bs, PetscSF *newsf) { - PetscInt i; auto n = sf.size(); PetscFunctionBegin; @@ -159,7 +158,7 @@ PetscErrorCode CreateCombinedSF(PC pc, const std::vector& sf, const std * allRoots: number of owned global dofs; * allLeaves: number of visible dofs (global + ghosted). */ - for (i = 0; i < n; ++i) { + for (size_t i = 0; i < n; ++i) { PetscInt nroots, nleaves; PetscCall(PetscSFGetGraph(sf[i], &nroots, &nleaves, NULL, NULL)); @@ -170,7 +169,7 @@ PetscErrorCode CreateCombinedSF(PC pc, const std::vector& sf, const std PetscCall(PetscMalloc1(allLeaves, &iremote)); // Now build an SF that just contains process connectivity. PetscCall(PetscHSetICreate(&ranksUniq)); - for (i = 0; i < n; ++i) { + for (size_t i = 0; i < n; ++i) { const PetscMPIInt *ranks = NULL; PetscMPIInt nranks, j; @@ -187,7 +186,7 @@ PetscErrorCode CreateCombinedSF(PC pc, const std::vector& sf, const std PetscCall(PetscHSetIGetElems(ranksUniq, &index, ranks)); PetscCall(PetscHMapICreate(&rankToIndex)); - for (i = 0; i < numRanks; ++i) { + for (PetscInt i = 0; i < numRanks; ++i) { remote[i].rank = ranks[i]; remote[i].index = 0; PetscCall(PetscHMapISet(rankToIndex, ranks[i], i)); @@ -203,7 +202,7 @@ PetscErrorCode CreateCombinedSF(PC pc, const std::vector& sf, const std PetscCall(PetscMalloc1(n*numRanks, &remoteOffsets)); offsets[0] = 0; - for (i = 1; i < n; ++i) { + for (size_t i = 1; i < n; ++i) { PetscInt nroots; PetscCall(PetscSFGetGraph(sf[i-1], &nroots, NULL, NULL, NULL)); @@ -211,8 +210,8 @@ PetscErrorCode CreateCombinedSF(PC pc, const std::vector& sf, const std } /* Offsets are the offsets on the current process of the * global dof numbering for the subspaces. */ - PetscCall(MPI_Type_contiguous(n, MPIU_INT, &contig)); - PetscCall(MPI_Type_commit(&contig)); + PetscCallMPI(MPI_Type_contiguous(n, MPIU_INT, &contig)); + PetscCallMPI(MPI_Type_commit(&contig)); #if MY_PETSC_VERSION_LT(3, 14, 4) PetscCall(PetscSFBcastBegin(rankSF, contig, offsets, remoteOffsets)); @@ -221,14 +220,14 @@ PetscErrorCode CreateCombinedSF(PC pc, const std::vector& sf, const std PetscCall(PetscSFBcastBegin(rankSF, contig, offsets, remoteOffsets, MPI_REPLACE)); PetscCall(PetscSFBcastEnd(rankSF, contig, offsets, remoteOffsets, MPI_REPLACE)); #endif - PetscCall(MPI_Type_free(&contig)); + PetscCallMPI(MPI_Type_free(&contig)); PetscCall(PetscFree(offsets)); PetscCall(PetscSFDestroy(&rankSF)); /* Now remoteOffsets contains the offsets on the remote * processes who communicate with me. So now we can * concatenate the list of SFs into a single one. */ index = 0; - for (i = 0; i < n; ++i) { + for (size_t i = 0; i < n; ++i) { const PetscSFNode *remote = NULL; const PetscInt *local = NULL; PetscInt nroots, nleaves, j; @@ -256,7 +255,7 @@ PetscErrorCode CreateCombinedSF(PC pc, const std::vector& sf, const std PetscCall(PetscSFCreate(PetscObjectComm((PetscObject)pc), newsf)); PetscCall(PetscSFSetGraph(*newsf, allRoots, allLeaves, ilocal, PETSC_OWN_POINTER, iremote, PETSC_OWN_POINTER)); } - PetscFunctionReturn(0); + PetscFunctionReturn(PETSC_SUCCESS); } @@ -266,7 +265,7 @@ PetscErrorCode PCSetup_TinyASM(PC pc) { auto blockjacobi = (BlockJacobi *)pc->data; blockjacobi -> updateValuesPerBlock(P); PetscCall(PetscLogEventEnd(PC_tinyasm_setup, pc, 0, 0, 0)); - return 0; + PetscFunctionReturn(PETSC_SUCCESS); } PetscErrorCode PCApply_TinyASM(PC pc, Vec b, Vec x) { @@ -295,13 +294,13 @@ PetscErrorCode PCApply_TinyASM(PC pc, Vec b, Vec x) { PetscCall(PetscSFReduceEnd(blockjacobi->sf, MPIU_SCALAR, &(blockjacobi->localx[0]), globalx, MPI_SUM)); PetscCall(VecRestoreArray(x, &globalx)); PetscCall(PetscLogEventEnd(PC_tinyasm_apply, pc, 0, 0, 0)); - return 0; + PetscFunctionReturn(PETSC_SUCCESS); } PetscErrorCode PCDestroy_TinyASM(PC pc) { if(pc->data) delete (BlockJacobi *)pc->data; - return 0; + PetscFunctionReturn(PETSC_SUCCESS); } PetscErrorCode PCView_TinyASM(PC pc, PetscViewer viewer) { @@ -320,7 +319,7 @@ PetscErrorCode PCView_TinyASM(PC pc, PetscViewer viewer) { PetscCall(PetscViewerASCIIPrintf(viewer, "Largest block size %" PetscInt_FMT " \n", biggestblock)); PetscCall(PetscViewerASCIIPopTab(viewer)); } - return 0; + PetscFunctionReturn(PETSC_SUCCESS); } PetscErrorCode PCCreate_TinyASM(PC pc) { @@ -329,7 +328,7 @@ PetscErrorCode PCCreate_TinyASM(PC pc) { pc->ops->setup = PCSetup_TinyASM; pc->ops->destroy = PCDestroy_TinyASM; pc->ops->view = PCView_TinyASM; - return 0; + PetscFunctionReturn(PETSC_SUCCESS); } // pybind11 casters for PETSc/petsc4py objects, copied from dolfinx repo // Import petsc4py on demand @@ -385,12 +384,15 @@ PYBIND11_MODULE(_tinyasm, m) { PetscLogEventRegister("PCTinyASMApply", PC_CLASSID, &PC_tinyasm_apply); m.def("SetASMLocalSubdomains", [](PC pc, std::vector ises, std::vector sfs, std::vector blocksizes, int localsize) { - PetscInt i, p, numDofs; - PetscCall(PetscLogEventBegin(PC_tinyasm_SetASMLocalSubdomains, pc, 0, 0, 0)); + PetscInt p, numDofs; + + MPI_Comm comm = PetscObjectComm((PetscObject) pc); + + PetscCallAbort(comm, PetscLogEventBegin(PC_tinyasm_SetASMLocalSubdomains, pc, 0, 0, 0)); auto P = pc->pmat; ISLocalToGlobalMapping lgr; ISLocalToGlobalMapping lgc; - MatGetLocalToGlobalMapping(P, &lgr, &lgc); + PetscCallAbort(comm, MatGetLocalToGlobalMapping(P, &lgr, &lgc)); int numBlocks = ises.size(); vector> dofsPerBlock(numBlocks); @@ -398,27 +400,27 @@ PYBIND11_MODULE(_tinyasm, m) { const PetscInt* isarray; for (p = 0; p < numBlocks; p++) { - PetscCall(ISGetSize(ises[p], &numDofs)); - PetscCall(ISGetIndices(ises[p], &isarray)); + PetscCallAbort(comm, ISGetSize(ises[p], &numDofs)); + PetscCallAbort(comm, ISGetIndices(ises[p], &isarray)); dofsPerBlock[p] = vector(); dofsPerBlock[p].reserve(numDofs); globalDofsPerBlock[p] = vector(numDofs, 0); - for (i = 0; i < numDofs; i++) { + for (PetscInt i = 0; i < numDofs; i++) { dofsPerBlock[p].push_back(isarray[i]); } - PetscCall(ISRestoreIndices(ises[p], &isarray)); - ISLocalToGlobalMappingApply(lgr, numDofs, &dofsPerBlock[p][0], &globalDofsPerBlock[p][0]); + PetscCallAbort(comm, ISRestoreIndices(ises[p], &isarray)); + PetscCallAbort(comm, ISLocalToGlobalMappingApply(lgr, numDofs, &dofsPerBlock[p][0], &globalDofsPerBlock[p][0])); } DM dm; - PetscCall(PCGetDM(pc, &dm)); + PetscCallAbort(comm, PCGetDM(pc, &dm)); PetscSF newsf; - PetscCall(CreateCombinedSF(pc, sfs, blocksizes, &newsf)); + PetscCallAbort(comm, CreateCombinedSF(pc, sfs, blocksizes, &newsf)); auto blockjacobi = new BlockJacobi(dofsPerBlock, globalDofsPerBlock, localsize, newsf); pc->data = (void*)blockjacobi; - PetscCall(PetscLogEventEnd(PC_tinyasm_SetASMLocalSubdomains, pc, 0, 0, 0)); + PetscCallAbort(comm, PetscLogEventEnd(PC_tinyasm_SetASMLocalSubdomains, pc, 0, 0, 0)); return 0; }); } From ff55c4cdac082c3b513fbb3bc5b81871c0476faa Mon Sep 17 00:00:00 2001 From: Robert Kirby Date: Mon, 13 Jan 2025 17:17:22 -0600 Subject: [PATCH 05/10] extra / (#3970) --- docs/source/firedrake_usa_25.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/firedrake_usa_25.rst b/docs/source/firedrake_usa_25.rst index 3bc4cb4381..e9cf3e7d6a 100644 --- a/docs/source/firedrake_usa_25.rst +++ b/docs/source/firedrake_usa_25.rst @@ -46,7 +46,7 @@ Registration ------------ Register for the conference at -`this link `__. +`this link `__. The registration fees are as follows: From 723ca723657332963b4f583a2772e44b07ce1781 Mon Sep 17 00:00:00 2001 From: Robert Kirby Date: Tue, 14 Jan 2025 14:32:34 -0600 Subject: [PATCH 06/10] add link to student support form (#3974) --- docs/source/firedrake_usa_25.rst | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/source/firedrake_usa_25.rst b/docs/source/firedrake_usa_25.rst index e9cf3e7d6a..0f5f617956 100644 --- a/docs/source/firedrake_usa_25.rst +++ b/docs/source/firedrake_usa_25.rst @@ -60,6 +60,8 @@ The registration fees are as follows: - $200 The `SIAM Texas-Louisiana Section `__ is providing some support for students currently attending universities in Texas or Louisiana to attend. +Students may submit an application to be considered for funding `here `__. + Abstract submission From cbfc2b2f359a2524b0871ead22e6b297fe8929ff Mon Sep 17 00:00:00 2001 From: MohamedAlyLoutfy <73834061+MohamedAlyLoutfy@users.noreply.github.com> Date: Wed, 15 Jan 2025 16:34:52 +0100 Subject: [PATCH 07/10] Preserve default mesh name after Netgen mesh creation (#3951) * Preserve default mesh name after Netgen mesh creation to avoid checkpointing errors * Ensure consistent handling of mesh and topology names for Netgen and non-Netgen meshes --- firedrake/mesh.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/firedrake/mesh.py b/firedrake/mesh.py index 7e9a61773b..8a0f8ec8a9 100644 --- a/firedrake/mesh.py +++ b/firedrake/mesh.py @@ -3073,6 +3073,7 @@ def Mesh(meshfile, **kwargs): netgen_flags = kwargs.get("netgen_flags", {"quad": False, "transform": None, "purify_to_tets": False}) netgen_firedrake_mesh = FiredrakeMesh(meshfile, netgen_flags, user_comm) plex = netgen_firedrake_mesh.meshMap.petscPlex + plex.setName(_generate_default_mesh_topology_name(name)) else: basename, ext = os.path.splitext(meshfile) if ext.lower() in ['.e', '.exo']: @@ -3099,10 +3100,11 @@ def Mesh(meshfile, **kwargs): distribution_name=kwargs.get("distribution_name"), permutation_name=kwargs.get("permutation_name"), comm=user_comm) - mesh = make_mesh_from_mesh_topology(topology, name) if netgen and isinstance(meshfile, netgen.libngpy._meshing.Mesh): - netgen_firedrake_mesh.createFromTopology(topology, name=plex.getName(), comm=user_comm) + netgen_firedrake_mesh.createFromTopology(topology, name=name, comm=user_comm) mesh = netgen_firedrake_mesh.firedrakeMesh + else: + mesh = make_mesh_from_mesh_topology(topology, name) mesh._tolerance = tolerance return mesh From cd68f9dae50998cc9bdbf50fa0bb9972c919d3b3 Mon Sep 17 00:00:00 2001 From: "Patrick E. Farrell" Date: Wed, 15 Jan 2025 16:19:49 +0000 Subject: [PATCH 08/10] Add functionality to print patch statistics in ASMStarPC and friends (#3875) * Add functionality to print patch statistics in ASMStarPC and friends --------- Co-authored-by: Pablo Brubeck --- firedrake/preconditioners/asm.py | 20 +++++++++++++++++++- tests/firedrake/regression/test_star_pc.py | 4 +++- 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/firedrake/preconditioners/asm.py b/firedrake/preconditioners/asm.py index 19738fe438..f3eaa037d9 100644 --- a/firedrake/preconditioners/asm.py +++ b/firedrake/preconditioners/asm.py @@ -6,6 +6,7 @@ from firedrake.dmhooks import get_function_space from firedrake.logging import warning from tinyasm import _tinyasm as tinyasm +from mpi4py import MPI import numpy @@ -33,7 +34,7 @@ def initialize(self, pc): # Extract function space and mesh to obtain plex and indexing functions V = get_function_space(dm) - # Obtain patches from user defined funtion + # Obtain patches from user defined function ises = self.get_patches(V) # PCASM expects at least one patch, so we define an empty one on idle processes if len(ises) == 0: @@ -90,6 +91,20 @@ def initialize(self, pc): asmpc.setFromOptions() self.asmpc = asmpc + self._patch_statistics = [] + if opts.getBool("view_patch_sizes", default=False): + # Compute and stash patch statistics + mpi_comm = pc.comm.tompi4py() + max_local_patch = max(is_.getSize() for is_ in ises) + min_local_patch = min(is_.getSize() for is_ in ises) + sum_local_patch = sum(is_.getSize() for is_ in ises) + max_global_patch = mpi_comm.allreduce(max_local_patch, op=MPI.MAX) + min_global_patch = mpi_comm.allreduce(min_local_patch, op=MPI.MIN) + sum_global_patch = mpi_comm.allreduce(sum_local_patch, op=MPI.SUM) + avg_global_patch = sum_global_patch / mpi_comm.allreduce(len(ises) if sum_local_patch > 0 else 0, op=MPI.SUM) + msg = f"Minimum / average / maximum patch sizes : {min_global_patch} / {avg_global_patch} / {max_global_patch}\n" + self._patch_statistics.append(msg) + @abc.abstractmethod def get_patches(self, V): ''' Get the patches used for PETSc PCASM @@ -103,6 +118,9 @@ def get_patches(self, V): def view(self, pc, viewer=None): self.asmpc.view(viewer=viewer) + if viewer is not None: + for msg in self._patch_statistics: + viewer.printfASCII(msg) def update(self, pc): # This is required to update an inplace ILU factorization diff --git a/tests/firedrake/regression/test_star_pc.py b/tests/firedrake/regression/test_star_pc.py index ba1f065eca..eb02c665e8 100644 --- a/tests/firedrake/regression/test_star_pc.py +++ b/tests/firedrake/regression/test_star_pc.py @@ -41,6 +41,7 @@ def test_star_equivalence(problem_type, backend): star_params = {"mat_type": "aij", "snes_type": "ksponly", + "snes_view": None, "ksp_type": "richardson", "pc_type": "mg", "pc_mg_type": "multiplicative", @@ -50,7 +51,8 @@ def test_star_equivalence(problem_type, backend): "mg_levels_ksp_max_it": 1, "mg_levels_pc_type": "python", "mg_levels_pc_python_type": "firedrake.ASMStarPC", - "mg_levels_pc_star_construct_dim": 0} + "mg_levels_pc_star_construct_dim": 0, + "mg_levels_pc_star_view_patch_sizes": None} comp_params = {"mat_type": "aij", "snes_type": "ksponly", From ad9fe2c762253b2317f5d24d68ae56e19b944b4e Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 15 Jan 2025 16:26:02 +0000 Subject: [PATCH 09/10] Test high order CR and variants (#3969) --- tests/firedrake/regression/test_helmholtz.py | 2 +- .../test_helmholtz_crouzeix_raviart.py | 59 +++++++++++++++++++ .../regression/test_helmholtz_serendipity.py | 7 +-- 3 files changed, 62 insertions(+), 6 deletions(-) create mode 100644 tests/firedrake/regression/test_helmholtz_crouzeix_raviart.py diff --git a/tests/firedrake/regression/test_helmholtz.py b/tests/firedrake/regression/test_helmholtz.py index 9097840275..6eeba6e71a 100644 --- a/tests/firedrake/regression/test_helmholtz.py +++ b/tests/firedrake/regression/test_helmholtz.py @@ -49,7 +49,7 @@ def helmholtz(r, quadrilateral=False, degree=2, mesh=None): assemble(L) sol = Function(V) solve(a == L, sol, solver_parameters={'ksp_type': 'cg'}) - # Analytical solution + # Error norm return sqrt(assemble(inner(sol - expect, sol - expect) * dx)), sol, expect diff --git a/tests/firedrake/regression/test_helmholtz_crouzeix_raviart.py b/tests/firedrake/regression/test_helmholtz_crouzeix_raviart.py new file mode 100644 index 0000000000..9c8faef8a3 --- /dev/null +++ b/tests/firedrake/regression/test_helmholtz_crouzeix_raviart.py @@ -0,0 +1,59 @@ +"""This demo program solves Helmholtz's equation + + - div grad u(x, y) + u(x,y) = f(x, y) + +on the unit square with source f given by + + f(x, y) = (1.0 + 8.0*pi**2)*cos(x[0]*2*pi)*cos(x[1]*2*pi) + +and the analytical solution + + u(x, y) = cos(x[0]*2*pi)*cos(x[1]*2*pi) +""" + +import numpy as np +import pytest + +from firedrake import * + + +def helmholtz(r, quadrilateral=False, degree=1, variant=None, mesh=None): + # Create mesh and define function space + if mesh is None: + mesh = UnitSquareMesh(2 ** r, 2 ** r, quadrilateral=quadrilateral) + V = FunctionSpace(mesh, "CR", degree, variant=variant) + + x, y = SpatialCoordinate(mesh) + + # Define variational problem + u = TrialFunction(V) + v = TestFunction(V) + + uex = cos(x*pi*2)*cos(y*pi*2) + f = -div(grad(uex)) + uex + + a = (inner(grad(u), grad(v)) + inner(u, v))*dx + L = inner(f, v)*dx(degree=12) + + params = {"snes_type": "ksponly", + "ksp_type": "preonly", + "pc_type": "lu"} + + # Compute solution + sol = Function(V) + solve(a == L, sol, solver_parameters=params) + # Error norm + return sqrt(assemble(dot(sol - uex, sol - uex) * dx)), sol, uex + + +@pytest.mark.parametrize(('testcase', 'convrate'), + [((1, (4, 6)), 1.9), + ((3, (2, 4)), 3.9), + ((5, (2, 4)), 5.7)]) +@pytest.mark.parametrize("variant", ("point", "integral")) +def test_firedrake_helmholtz_scalar_convergence(variant, testcase, convrate): + degree, (start, end) = testcase + l2err = np.zeros(end - start) + for ii in [i + start for i in range(len(l2err))]: + l2err[ii - start] = helmholtz(ii, degree=degree, variant=variant)[0] + assert (np.array([np.log2(l2err[i]/l2err[i+1]) for i in range(len(l2err)-1)]) > convrate).all() diff --git a/tests/firedrake/regression/test_helmholtz_serendipity.py b/tests/firedrake/regression/test_helmholtz_serendipity.py index 77578d3c86..937b02899b 100644 --- a/tests/firedrake/regression/test_helmholtz_serendipity.py +++ b/tests/firedrake/regression/test_helmholtz_serendipity.py @@ -35,7 +35,7 @@ def helmholtz(r, quadrilateral=True, degree=2, mesh=None): uex = cos(x*pi*2)*cos(y*pi*2) f = -div(grad(uex)) + uex - a = (inner(grad(u), grad(v)) + inner(u, v))*dx(degree=12) + a = (inner(grad(u), grad(v)) + inner(u, v))*dx L = inner(f, v)*dx(degree=12) params = {"snes_type": "ksponly", @@ -45,10 +45,7 @@ def helmholtz(r, quadrilateral=True, degree=2, mesh=None): # Compute solution sol = Function(V) solve(a == L, sol, solver_parameters=params) - - # Analytical solution - f = Function(V) - f.project(cos(x*pi*2)*cos(y*pi*2)) + # Error norm return sqrt(assemble(dot(sol - uex, sol - uex) * dx)), sol, uex From 9c5ec2fc55410a4e9e6c1e226207cddd54a3ff1f Mon Sep 17 00:00:00 2001 From: Pablo Brubeck Date: Wed, 15 Jan 2025 22:09:19 +0000 Subject: [PATCH 10/10] Fieldsplit: replace empty Forms with ZeroBaseForm (#3947) * Restricted Cofunction RHS * Fix BCs on Cofunction * LinearSolver: check function spaces * assemble(form, zero_bc_nodes=True) as default * Fix FunctionAssignBlock * Allow Cofunction.assign take in constants * remove BaseFormAssembler test * only supply relevant kwargs to OneFormAssembler * Only interpolate the residual, not every cofunction in the RHS * Fix tests * Fix adjoint utils * More robust test for (unrestricted) Cofunction RHS * Replace empty Jacobians with ZeroBaseForm * Do not split off-diagonal blocks if we only want the diagonal * Zero-simplify slate Tensors * set bcs directly on diagonal Cofunction * ImplicitMatrixContext: handle empty action * Only extract constants referenced in the kernel * Adjoint: only skip expand_derivatives if necessary * EquationBC: do not reconstruct empty Forms * lower degree for EquationBC tests * Update .github/workflows/build.yml --- demos/netgen/netgen_mesh.py.rst | 3 +- .../adjoint_utils/blocks/dirichlet_bc.py | 6 +- firedrake/adjoint_utils/blocks/function.py | 1 + firedrake/adjoint_utils/blocks/solving.py | 26 ++-- firedrake/adjoint_utils/variational_solver.py | 13 +- firedrake/assemble.py | 75 ++++++----- firedrake/bcs.py | 8 +- firedrake/cofunction.py | 6 +- firedrake/formmanipulation.py | 126 +++++++----------- firedrake/functionspaceimpl.py | 11 +- firedrake/linear_solver.py | 7 + firedrake/matrix_free/operators.py | 34 ++--- firedrake/preconditioners/massinv.py | 2 +- firedrake/slate/slate.py | 90 ++++++++----- firedrake/slate/static_condensation/scpc.py | 2 +- firedrake/solving_utils.py | 13 +- firedrake/tsfc_interface.py | 4 +- firedrake/variational_solver.py | 12 +- .../equation_bcs/test_equation_bcs.py | 81 ++++++----- tests/firedrake/multigrid/test_poisson_gmg.py | 7 +- tests/firedrake/regression/test_assemble.py | 2 +- .../regression/test_assemble_baseform.py | 47 +------ tests/firedrake/regression/test_bcs.py | 2 +- tests/firedrake/regression/test_netgen.py | 6 +- .../test_restricted_function_space.py | 9 +- .../regression/test_solving_interface.py | 20 +-- .../firedrake/slate/test_assemble_tensors.py | 17 ++- 27 files changed, 310 insertions(+), 320 deletions(-) diff --git a/demos/netgen/netgen_mesh.py.rst b/demos/netgen/netgen_mesh.py.rst index 47fe769e70..3947cc828b 100755 --- a/demos/netgen/netgen_mesh.py.rst +++ b/demos/netgen/netgen_mesh.py.rst @@ -380,8 +380,7 @@ We will now show how to solve the Poisson problem on a high-order mesh, of order bc = DirichletBC(V, 0.0, [1]) A = assemble(a, bcs=bc) - b = assemble(l) - bc.apply(b) + b = assemble(l, bcs=bc) solve(A, sol, b, solver_parameters={"ksp_type": "cg", "pc_type": "lu"}) VTKFile("output/Sphere.pvd").write(sol) diff --git a/firedrake/adjoint_utils/blocks/dirichlet_bc.py b/firedrake/adjoint_utils/blocks/dirichlet_bc.py index b06d367da1..a918fb8a9e 100644 --- a/firedrake/adjoint_utils/blocks/dirichlet_bc.py +++ b/firedrake/adjoint_utils/blocks/dirichlet_bc.py @@ -51,7 +51,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, adj_output = None for adj_input in adj_inputs: if isconstant(c): - adj_value = firedrake.Function(self.parent_space.dual()) + adj_value = firedrake.Function(self.parent_space) adj_input.apply(adj_value) if self.function_space != self.parent_space: vec = extract_bc_subvector( @@ -88,11 +88,11 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, # you can even use the Function outside its domain. # For now we will just assume the FunctionSpace is the same for # the BC and the Function. - adj_value = firedrake.Function(self.parent_space.dual()) + adj_value = firedrake.Function(self.parent_space) adj_input.apply(adj_value) r = extract_bc_subvector( adj_value, c.function_space(), bc - ) + ).riesz_representation("l2") if adj_output is None: adj_output = r else: diff --git a/firedrake/adjoint_utils/blocks/function.py b/firedrake/adjoint_utils/blocks/function.py index dcb02da108..ae324fe1ad 100644 --- a/firedrake/adjoint_utils/blocks/function.py +++ b/firedrake/adjoint_utils/blocks/function.py @@ -79,6 +79,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, ) diff_expr_assembled = firedrake.Function(adj_input_func.function_space()) diff_expr_assembled.interpolate(ufl.conj(diff_expr)) + diff_expr_assembled = diff_expr_assembled.riesz_representation(riesz_map="l2") adj_output = firedrake.Function( R, val=firedrake.assemble(ufl.Action(diff_expr_assembled, adj_input_func)) ) diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index e4664665b0..2cf6d9fd36 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -197,14 +197,12 @@ def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs): def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy): dJdu_copy = dJdu.copy() - kwargs = self.assemble_kwargs.copy() # Homogenize and apply boundary conditions on adj_dFdu and dJdu. bcs = self._homogenize_bcs() - kwargs["bcs"] = bcs - dFdu = self._assemble_dFdu_adj(dFdu_adj_form, **kwargs) + dFdu = firedrake.assemble(dFdu_adj_form, bcs=bcs, **self.assemble_kwargs) for bc in bcs: - bc.apply(dJdu) + bc.zero(dJdu) adj_sol = firedrake.Function(self.function_space) firedrake.solve( @@ -219,10 +217,8 @@ def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy): return adj_sol, adj_sol_bdy def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu): - adj_sol_bdy = firedrake.Function( - self.function_space.dual(), dJdu.dat - firedrake.assemble( - firedrake.action(dFdu_adj_form, adj_sol)).dat) - return adj_sol_bdy + adj_sol_bdy = firedrake.assemble(dJdu - firedrake.action(dFdu_adj_form, adj_sol)) + return adj_sol_bdy.riesz_representation("l2") def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): @@ -264,8 +260,11 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, return dFdm dFdm = -firedrake.derivative(F_form, c_rep, trial_function) - dFdm = firedrake.adjoint(dFdm) - dFdm = dFdm * adj_sol + if isinstance(dFdm, ufl.Form): + dFdm = firedrake.adjoint(dFdm) + dFdm = firedrake.action(dFdm, adj_sol) + else: + dFdm = dFdm(adj_sol) dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) return dFdm @@ -654,9 +653,8 @@ def _forward_solve(self, lhs, rhs, func, bcs, **kwargs): def _adjoint_solve(self, dJdu, compute_bdy): dJdu_copy = dJdu.copy() # Homogenize and apply boundary conditions on adj_dFdu and dJdu. - bcs = self._homogenize_bcs() - for bc in bcs: - bc.apply(dJdu) + for bc in self.bcs: + bc.zero(dJdu) if ( self._ad_solvers["forward_nlvs"]._problem._constant_jacobian @@ -876,7 +874,7 @@ def __init__(self, source, target_space, target, bcs=[], **kwargs): self.add_dependency(bc, no_duplicates=True) def apply_mixedmass(self, a): - b = firedrake.Function(self.target_space) + b = firedrake.Function(self.target_space.dual()) with a.dat.vec_ro as vsrc, b.dat.vec_wo as vrhs: self.mixed_mass.mult(vsrc, vrhs) return b diff --git a/firedrake/adjoint_utils/variational_solver.py b/firedrake/adjoint_utils/variational_solver.py index c90d2668e0..c191308adc 100644 --- a/firedrake/adjoint_utils/variational_solver.py +++ b/firedrake/adjoint_utils/variational_solver.py @@ -2,6 +2,7 @@ from functools import wraps from pyadjoint.tape import get_working_tape, stop_annotating, annotate_tape, no_annotations from firedrake.adjoint_utils.blocks import NonlinearVariationalSolveBlock +from firedrake.ufl_expr import derivative, adjoint from ufl import replace @@ -11,7 +12,6 @@ def _ad_annotate_init(init): @no_annotations @wraps(init) def wrapper(self, *args, **kwargs): - from firedrake import derivative, adjoint, TrialFunction init(self, *args, **kwargs) self._ad_F = self.F self._ad_u = self.u_restrict @@ -20,10 +20,13 @@ def wrapper(self, *args, **kwargs): try: # Some forms (e.g. SLATE tensors) are not currently # differentiable. - dFdu = derivative(self.F, - self.u_restrict, - TrialFunction(self.u_restrict.function_space())) - self._ad_adj_F = adjoint(dFdu) + dFdu = derivative(self.F, self.u_restrict) + try: + self._ad_adj_F = adjoint(dFdu) + except ValueError: + # Try again without expanding derivatives, + # as dFdu might have been simplied to an empty Form + self._ad_adj_F = adjoint(dFdu, derivatives_expanded=True) except (TypeError, NotImplementedError): self._ad_adj_F = None self._ad_kwargs = {'Jp': self.Jp, 'form_compiler_parameters': self.form_compiler_parameters, 'is_linear': self.is_linear} diff --git a/firedrake/assemble.py b/firedrake/assemble.py index f3049ae01c..d08e64a487 100644 --- a/firedrake/assemble.py +++ b/firedrake/assemble.py @@ -83,7 +83,7 @@ def assemble(expr, *args, **kwargs): zero_bc_nodes : bool If `True`, set the boundary condition nodes in the output tensor to zero rather than to the values prescribed by the - boundary condition. Default is `False`. + boundary condition. Default is `True`. diagonal : bool If assembling a matrix is it diagonal? weight : float @@ -143,7 +143,6 @@ def get_assembler(form, *args, **kwargs): """ is_base_form_preprocessed = kwargs.pop('is_base_form_preprocessed', False) - bcs = kwargs.get('bcs', None) fc_params = kwargs.get('form_compiler_parameters', None) if isinstance(form, ufl.form.BaseForm) and not is_base_form_preprocessed: mat_type = kwargs.get('mat_type', None) @@ -155,8 +154,13 @@ def get_assembler(form, *args, **kwargs): if len(form.arguments()) == 0: return ZeroFormAssembler(form, form_compiler_parameters=fc_params) elif len(form.arguments()) == 1 or diagonal: - return OneFormAssembler(form, *args, bcs=bcs, form_compiler_parameters=fc_params, needs_zeroing=kwargs.get('needs_zeroing', True), - zero_bc_nodes=kwargs.get('zero_bc_nodes', False), diagonal=diagonal) + return OneFormAssembler(form, *args, + bcs=kwargs.get("bcs", None), + form_compiler_parameters=fc_params, + needs_zeroing=kwargs.get("needs_zeroing", True), + zero_bc_nodes=kwargs.get("zero_bc_nodes", True), + diagonal=diagonal, + weight=kwargs.get("weight", 1.0)) elif len(form.arguments()) == 2: return TwoFormAssembler(form, *args, **kwargs) else: @@ -308,7 +312,7 @@ def __init__(self, sub_mat_type=None, options_prefix=None, appctx=None, - zero_bc_nodes=False, + zero_bc_nodes=True, diagonal=False, weight=1.0, allocation_integral_types=None): @@ -381,6 +385,12 @@ def visitor(e, *operands): visited = {} result = BaseFormAssembler.base_form_postorder_traversal(self._form, visitor, visited) + # Apply BCs after assembly + rank = len(self._form.arguments()) + if rank == 1 and not isinstance(result, ufl.ZeroBaseForm): + for bc in self._bcs: + bc.zero(result) + if tensor: BaseFormAssembler.update_tensor(result, tensor) return tensor @@ -405,8 +415,8 @@ def base_form_assembly_visitor(self, expr, tensor, *args): if rank == 0: assembler = ZeroFormAssembler(form, form_compiler_parameters=self._form_compiler_params) elif rank == 1 or (rank == 2 and self._diagonal): - assembler = OneFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, - zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal) + assembler = OneFormAssembler(form, form_compiler_parameters=self._form_compiler_params, + zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal, weight=self._weight) elif rank == 2: assembler = TwoFormAssembler(form, bcs=self._bcs, form_compiler_parameters=self._form_compiler_params, mat_type=self._mat_type, sub_mat_type=self._sub_mat_type, @@ -577,10 +587,15 @@ def base_form_assembly_visitor(self, expr, tensor, *args): @staticmethod def update_tensor(assembled_base_form, tensor): if isinstance(tensor, (firedrake.Function, firedrake.Cofunction)): - assembled_base_form.dat.copy(tensor.dat) + if isinstance(assembled_base_form, ufl.ZeroBaseForm): + tensor.dat.zero() + else: + assembled_base_form.dat.copy(tensor.dat) elif isinstance(tensor, matrix.MatrixBase): - # Uses the PETSc copy method. - assembled_base_form.petscmat.copy(tensor.petscmat) + if isinstance(assembled_base_form, ufl.ZeroBaseForm): + tensor.petscmat.zeroEntries() + else: + assembled_base_form.petscmat.copy(tensor.petscmat) else: raise NotImplementedError("Cannot update tensor of type %s" % type(tensor)) @@ -807,9 +822,9 @@ def restructure_base_form(expr, visited=None): return ufl.action(expr, ustar) # -- Case (6) -- # - if isinstance(expr, ufl.FormSum) and all(isinstance(c, ufl.core.base_form_operator.BaseFormOperator) for c in expr.components()): - # Return ufl.Sum - return sum([c for c in expr.components()]) + if isinstance(expr, ufl.FormSum) and all(ufl.duals.is_dual(a.function_space()) for a in expr.arguments()): + # Return ufl.Sum if we are assembling a FormSum with Coarguments (a primal expression) + return sum(w*c for w, c in zip(expr.weights(), expr.components())) return expr @staticmethod @@ -1138,7 +1153,7 @@ class OneFormAssembler(ParloopFormAssembler): Parameters ---------- - form : ufl.Form or slate.TensorBasehe + form : ufl.Form or slate.TensorBase 1-form. Notes @@ -1149,14 +1164,15 @@ class OneFormAssembler(ParloopFormAssembler): @classmethod def _cache_key(cls, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=False, diagonal=False): + zero_bc_nodes=True, diagonal=False, weight=1.0): bcs = solving._extract_bcs(bcs) - return tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, zero_bc_nodes, diagonal + return tuple(bcs), tuplify(form_compiler_parameters), needs_zeroing, zero_bc_nodes, diagonal, weight @FormAssembler._skip_if_initialised def __init__(self, form, bcs=None, form_compiler_parameters=None, needs_zeroing=True, - zero_bc_nodes=False, diagonal=False): + zero_bc_nodes=True, diagonal=False, weight=1.0): super().__init__(form, bcs=bcs, form_compiler_parameters=form_compiler_parameters, needs_zeroing=needs_zeroing) + self._weight = weight self._diagonal = diagonal self._zero_bc_nodes = zero_bc_nodes if self._diagonal and any(isinstance(bc, EquationBCSplit) for bc in self._bcs): @@ -1185,23 +1201,21 @@ def _apply_bc(self, tensor, bc): elif isinstance(bc, EquationBCSplit): bc.zero(tensor) type(self)(bc.f, bcs=bc.bcs, form_compiler_parameters=self._form_compiler_params, needs_zeroing=False, - zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal).assemble(tensor=tensor) + zero_bc_nodes=self._zero_bc_nodes, diagonal=self._diagonal, weight=self._weight).assemble(tensor=tensor) else: raise AssertionError def _apply_dirichlet_bc(self, tensor, bc): - if not self._zero_bc_nodes: - tensor_func = tensor.riesz_representation(riesz_map="l2") - if self._diagonal: - bc.set(tensor_func, 1) - else: - bc.apply(tensor_func) - tensor.assign(tensor_func.riesz_representation(riesz_map="l2")) + if self._diagonal: + bc.set(tensor, self._weight) + elif not self._zero_bc_nodes: + # NOTE this only works if tensor is a Function and not a Cofunction + bc.apply(tensor) else: bc.zero(tensor) def _check_tensor(self, tensor): - if tensor.function_space() != self._form.arguments()[0].function_space(): + if tensor.function_space() != self._form.arguments()[0].function_space().dual(): raise ValueError("Form's argument does not match provided result tensor") @staticmethod @@ -2127,14 +2141,13 @@ def iter_active_coefficients(form, kinfo): @staticmethod def iter_constants(form, kinfo): - """Yield the form constants""" + """Yield the form constants referenced in ``kinfo``.""" if isinstance(form, slate.TensorBase): - for const in form.constants(): - yield const + all_constants = form.constants() else: all_constants = extract_firedrake_constants(form) - for constant_index in kinfo.constant_numbers: - yield all_constants[constant_index] + for constant_index in kinfo.constant_numbers: + yield all_constants[constant_index] @staticmethod def index_function_spaces(form, indices): diff --git a/firedrake/bcs.py b/firedrake/bcs.py index f0d007ede4..5884907feb 100644 --- a/firedrake/bcs.py +++ b/firedrake/bcs.py @@ -634,10 +634,10 @@ def reconstruct(self, field=None, V=None, subu=None, u=None, row_field=None, col return rank = len(self.f.arguments()) splitter = ExtractSubBlock() - if rank == 1: - form = splitter.split(self.f, argument_indices=(row_field, )) - elif rank == 2: - form = splitter.split(self.f, argument_indices=(row_field, col_field)) + form = splitter.split(self.f, argument_indices=(row_field, col_field)[:rank]) + if isinstance(form, ufl.ZeroBaseForm) or form.empty(): + # form is empty, do nothing + return if u is not None: form = firedrake.replace(form, {self.u: u}) if action_x is not None: diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index ecfb386dd9..f46d7e0849 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -229,8 +229,10 @@ def assign(self, expr, subset=None, expr_from_assemble=False): return self.assign( assembled_expr, subset=subset, expr_from_assemble=True) - - raise ValueError('Cannot assign %s' % expr) + else: + from firedrake.assign import Assigner + Assigner(self, expr, subset).assign() + return self def riesz_representation(self, riesz_map='L2', **solver_options): """Return the Riesz representation of this :class:`Cofunction` with respect to the given Riesz map. diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index c3a3166ae8..1d2aa67aae 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -2,17 +2,27 @@ import numpy import collections -from ufl import as_vector, FormSum, Form, split +from ufl import as_vector, split from ufl.classes import Zero, FixedIndex, ListTensor, ZeroBaseForm from ufl.algorithms.map_integrands import map_integrand_dags +from ufl.algorithms import expand_derivatives from ufl.corealg.map_dag import MultiFunction, map_expr_dags from pyop2 import MixedDat +from pyop2.utils import as_tuple from firedrake.petsc import PETSc -from firedrake.ufl_expr import Argument +from firedrake.functionspace import MixedFunctionSpace from firedrake.cofunction import Cofunction -from firedrake.functionspace import FunctionSpace, MixedFunctionSpace, DualSpace + + +def subspace(V, indices): + """Construct a collapsed subspace using components from V.""" + if len(indices) == 1: + W = V[indices[0]] + else: + W = MixedFunctionSpace([V[i] for i in indices]) + return W.collapse() class ExtractSubBlock(MultiFunction): @@ -30,13 +40,19 @@ def indexed(self, o, child, multiindex): indices = multiindex.indices() if isinstance(child, ListTensor) and all(isinstance(i, FixedIndex) for i in indices): if len(indices) == 1: - return child.ufl_operands[indices[0]._value] + return child[indices[0]] + elif len(indices) == len(child.ufl_operands) and all(k == int(i) for k, i in enumerate(indices)): + return child else: - return ListTensor(*(child.ufl_operands[i._value] for i in multiindex.indices())) + return ListTensor(*(child[i] for i in indices)) return self.expr(o, child, multiindex) index_inliner = IndexInliner() + def _subspace_argument(self, a): + return type(a)(subspace(a.function_space(), self.blocks[a.number()]), + a.number(), part=a.part()) + @PETSc.Log.EventDecorator() def split(self, form, argument_indices): """Split a form. @@ -52,7 +68,7 @@ def split(self, form, argument_indices): """ args = form.arguments() self._arg_cache = {} - self.blocks = dict(enumerate(argument_indices)) + self.blocks = dict(enumerate(map(as_tuple, argument_indices))) if len(args) == 0: # Functional can't be split return form @@ -60,7 +76,11 @@ def split(self, form, argument_indices): assert (len(idx) == 1 for idx in self.blocks.values()) assert (idx[0] == 0 for idx in self.blocks.values()) return form + # TODO find a way to distinguish empty Forms avoiding expand_derivatives f = map_integrand_dags(self, form) + if expand_derivatives(f).empty(): + # Get ZeroBaseForm with the right shape + f = ZeroBaseForm(tuple(map(self._subspace_argument, form.arguments()))) return f expr = MultiFunction.reuse_if_untouched @@ -98,77 +118,34 @@ def argument(self, o): if o in self._arg_cache: return self._arg_cache[o] - V_is = V.subfunctions indices = self.blocks[o.number()] - # Only one index provided. - if isinstance(indices, int): - indices = (indices, ) + a = self._subspace_argument(o) + asplit = (a, ) if len(indices) == 1 else split(a) - if len(indices) == 1: - W = V_is[indices[0]] - W = FunctionSpace(W.mesh(), W.ufl_element()) - a = (Argument(W, o.number(), part=o.part()), ) - else: - W = MixedFunctionSpace([V_is[i] for i in indices]) - a = split(Argument(W, o.number(), part=o.part())) args = [] - for i in range(len(V_is)): + for i in range(len(V)): if i in indices: - c = indices.index(i) - a_ = a[c] - if len(a_.ufl_shape) == 0: - args += [a_] - else: - args += [a_[j] for j in numpy.ndindex(a_.ufl_shape)] + asub = asplit[indices.index(i)] + args.extend(asub[j] for j in numpy.ndindex(asub.ufl_shape)) else: - args += [Zero() - for j in numpy.ndindex(V_is[i].value_shape)] + args.extend(Zero() for j in numpy.ndindex(V[i].value_shape)) return self._arg_cache.setdefault(o, as_vector(args)) def cofunction(self, o): V = o.function_space() - # Not on a mixed space, just return ourselves. if len(V) == 1: + # Not on a mixed space, just return ourselves. return o # We only need the test space for Cofunction indices = self.blocks[0] - V_is = V.subfunctions - - # Only one index provided. - if isinstance(indices, int): - indices = (indices, ) - - # for two-forms, the cofunction should only - # be returned for the diagonal blocks, so - # if we are asked for an off-diagonal block - # then we return a zero form, analogously to - # the off components of arguments. - if len(self.blocks) == 2: - itest, itrial = self.blocks - on_diag = (itest == itrial) + W = subspace(V, indices) + if len(W) == 1: + return Cofunction(W, val=o.dat[indices[0]]) else: - on_diag = True - - # if we are on the diagonal, then return a Cofunction - # in the relevant subspace that points to the data in - # the full space. This means that the right hand side - # of the fieldsplit problem will be correct. - if on_diag: - if len(indices) == 1: - i = indices[0] - W = V_is[i] - W = DualSpace(W.mesh(), W.ufl_element()) - c = Cofunction(W, val=o.subfunctions[i].dat) - else: - W = MixedFunctionSpace([V_is[i] for i in indices]) - c = Cofunction(W, val=MixedDat(o.dat[i] for i in indices)) - else: - c = ZeroBaseForm(o.arguments()) - - return c + return Cofunction(W, val=MixedDat(o.dat[i] for i in indices)) SplitForm = collections.namedtuple("SplitForm", ["indices", "form"]) @@ -207,28 +184,15 @@ def split_form(form, diagonal=False): args = form.arguments() shape = tuple(len(a.function_space()) for a in args) forms = [] + rank = len(shape) if diagonal: - assert len(shape) == 2 + assert rank == 2 + rank = 1 for idx in numpy.ndindex(shape): + if diagonal: + i, j = idx + if i != j: + continue f = splitter.split(form, idx) - - # does f actually contain anything? - if isinstance(f, Cofunction): - flen = 1 - elif isinstance(f, FormSum): - flen = len(f.components()) - elif isinstance(f, Form): - flen = len(f.integrals()) - else: - raise ValueError( - "ExtractSubBlock.split should have returned an instance of " - "either Form, FormSum, or Cofunction") - - if flen > 0: - if diagonal: - i, j = idx - if i != j: - continue - idx = (i, ) - forms.append(SplitForm(indices=idx, form=f)) + forms.append(SplitForm(indices=idx[:rank], form=f)) return tuple(forms) diff --git a/firedrake/functionspaceimpl.py b/firedrake/functionspaceimpl.py index c366c03b66..f6a6f8e6b5 100644 --- a/firedrake/functionspaceimpl.py +++ b/firedrake/functionspaceimpl.py @@ -13,6 +13,7 @@ import ufl import finat.ufl +from ufl.duals import is_dual, is_primal from pyop2 import op2, mpi from pyop2.utils import as_tuple @@ -293,6 +294,9 @@ def restore_work_function(self, function): cache[function] = False def __eq__(self, other): + if is_primal(self) != is_primal(other) or \ + is_dual(self) != is_dual(other): + return False try: return self.topological == other.topological and \ self.mesh() is other.mesh() @@ -938,6 +942,9 @@ def __hash__(self): def local_to_global_map(self, bcs, lgmap=None): return lgmap or self.dof_dset.lgmap + def collapse(self): + return type(self)(self.function_space.collapse(), boundary_set=self.boundary_set) + class MixedFunctionSpace(object): r"""A function space on a mixed finite element. @@ -1230,16 +1237,16 @@ class ProxyRestrictedFunctionSpace(RestrictedFunctionSpace): r"""A :class:`RestrictedFunctionSpace` that one can attach extra properties to. :arg function_space: The function space to be restricted. - :kwarg name: The name of the restricted function space. :kwarg boundary_set: The boundary domains on which boundary conditions will be specified + :kwarg name: The name of the restricted function space. .. warning:: Users should not build a :class:`ProxyRestrictedFunctionSpace` directly, it is mostly used as an internal implementation detail. """ - def __new__(cls, function_space, name=None, boundary_set=frozenset()): + def __new__(cls, function_space, boundary_set=frozenset(), name=None): topology = function_space._mesh.topology self = super(ProxyRestrictedFunctionSpace, cls).__new__(cls) if function_space._mesh is not topology: diff --git a/firedrake/linear_solver.py b/firedrake/linear_solver.py index c1dfbcc07e..7721675a57 100644 --- a/firedrake/linear_solver.py +++ b/firedrake/linear_solver.py @@ -147,6 +147,13 @@ def solve(self, x, b): if not isinstance(b, (function.Function, cofunction.Cofunction)): raise TypeError("Provided RHS is a '%s', not a Function or Cofunction" % type(b).__name__) + # When solving `Ax = b`, with A: V x U -> R, or equivalently A: V -> U*, + # we need to make sure that x and b belong to V and U*, respectively. + if x.function_space() != self.trial_space: + raise ValueError(f"x must be a Function in {self.trial_space}.") + if b.function_space() != self.test_space.dual(): + raise ValueError(f"b must be a Cofunction in {self.test_space.dual()}.") + if len(self.trial_space) > 1 and self.nullspace is not None: self.nullspace._apply(self.trial_space.dof_dset.field_ises) if len(self.test_space) > 1 and self.transpose_nullspace is not None: diff --git a/firedrake/matrix_free/operators.py b/firedrake/matrix_free/operators.py index 3ee448730e..9a5a042f0e 100644 --- a/firedrake/matrix_free/operators.py +++ b/firedrake/matrix_free/operators.py @@ -10,6 +10,9 @@ from firedrake.bcs import DirichletBC, EquationBCSplit from firedrake.petsc import PETSc from firedrake.utils import cached_property +from firedrake.function import Function +from firedrake.cofunction import Cofunction +from ufl.form import ZeroBaseForm __all__ = ("ImplicitMatrixContext", ) @@ -107,23 +110,22 @@ def __init__(self, a, row_bcs=[], col_bcs=[], # create functions from test and trial space to help # with 1-form assembly - test_space, trial_space = [ - a.arguments()[i].function_space() for i in (0, 1) - ] - from firedrake import function, cofunction + test_space, trial_space = ( + arg.function_space() for arg in a.arguments() + ) # Need a cofunction since y receives the assembled result of Ax - self._ystar = cofunction.Cofunction(test_space.dual()) - self._y = function.Function(test_space) - self._x = function.Function(trial_space) - self._xstar = cofunction.Cofunction(trial_space.dual()) + self._ystar = Cofunction(test_space.dual()) + self._y = Function(test_space) + self._x = Function(trial_space) + self._xstar = Cofunction(trial_space.dual()) # These are temporary storage for holding the BC # values during matvec application. _xbc is for # the action and ._ybc is for transpose. if len(self.bcs) > 0: - self._xbc = cofunction.Cofunction(trial_space.dual()) + self._xbc = Cofunction(trial_space.dual()) if len(self.col_bcs) > 0: - self._ybc = cofunction.Cofunction(test_space.dual()) + self._ybc = Cofunction(test_space.dual()) # Get size information from template vecs on test and trial spaces trial_vec = trial_space.dof_dset.layout_vec @@ -135,6 +137,11 @@ def __init__(self, a, row_bcs=[], col_bcs=[], self.action = action(self.a, self._x) self.actionT = action(self.aT, self._y) + # TODO prevent action from returning empty Forms + if self.action.empty(): + self.action = ZeroBaseForm(self.a.arguments()[:-1]) + if self.actionT.empty(): + self.actionT = ZeroBaseForm(self.aT.arguments()[:-1]) # For assembling action(f, self._x) self.bcs_action = [] @@ -147,7 +154,7 @@ def __init__(self, a, row_bcs=[], col_bcs=[], self._assemble_action = get_assembler(self.action, bcs=self.bcs_action, form_compiler_parameters=self.fc_params, - zero_bc_nodes=True).assemble + ).assemble # For assembling action(adjoint(f), self._y) # Sorted list of equation bcs @@ -170,7 +177,6 @@ def __init__(self, a, row_bcs=[], col_bcs=[], @cached_property def _diagonal(self): - from firedrake import Cofunction assert self.on_diag return Cofunction(self._x.function_space().dual()) @@ -183,11 +189,9 @@ def _assemble_diagonal(self): def getDiagonal(self, mat, vec): self._assemble_diagonal(tensor=self._diagonal) - diagonal_func = self._diagonal.riesz_representation(riesz_map="l2") for bc in self.bcs: # Operator is identity on boundary nodes - bc.set(diagonal_func, 1) - self._diagonal.assign(diagonal_func.riesz_representation(riesz_map="l2")) + bc.set(self._diagonal, 1) with self._diagonal.dat.vec_ro as v: v.copy(vec) diff --git a/firedrake/preconditioners/massinv.py b/firedrake/preconditioners/massinv.py index 92f286c708..d29c704e8b 100644 --- a/firedrake/preconditioners/massinv.py +++ b/firedrake/preconditioners/massinv.py @@ -20,7 +20,7 @@ class MassInvPC(AssembledPC): context, keyed on ``"mu"``. """ def form(self, pc, test, trial): - _, bcs = super(MassInvPC, self).form(pc, test, trial) + _, bcs = super(MassInvPC, self).form(pc) appctx = self.get_appctx(pc) mu = appctx.get("mu", 1.0) diff --git a/firedrake/slate/slate.py b/firedrake/slate/slate.py index 85b7af3635..bb6909128f 100644 --- a/firedrake/slate/slate.py +++ b/firedrake/slate/slate.py @@ -21,7 +21,9 @@ from ufl import Constant from ufl.coefficient import BaseCoefficient +from firedrake.formmanipulation import ExtractSubBlock, subspace from firedrake.function import Function, Cofunction +from firedrake.ufl_expr import TestFunction from firedrake.utils import cached_property, unique from itertools import chain, count @@ -32,11 +34,9 @@ from ufl.corealg.multifunction import MultiFunction from ufl.classes import Zero from ufl.domain import join_domains, sort_domains -from ufl.form import Form +from ufl.form import BaseForm, Form, ZeroBaseForm import hashlib -from firedrake.formmanipulation import ExtractSubBlock - from tsfc.ufl_utils import extract_firedrake_constants @@ -237,7 +237,7 @@ def coeff_map(self): coeff_map[m].update(c.indices[0]) else: m = self.coefficients().index(c) - split_map = tuple(range(len(c.subfunctions))) if isinstance(c, Function) or isinstance(c, Constant) or isinstance(c, Cofunction) else tuple(range(1)) + split_map = tuple(range(len(c.subfunctions))) if isinstance(c, (Function, Constant, Cofunction)) else (0,) coeff_map[m].update(split_map) return tuple((k, tuple(sorted(v)))for k, v in coeff_map.items()) @@ -293,6 +293,10 @@ def solve(self, B, decomposition=None): """ return Solve(self, B, decomposition=decomposition) + def empty(self): + """Returns whether the form associated with the tensor is empty.""" + return False + @cached_property def blocks(self): """Returns an object containing the blocks of the tensor defined @@ -382,6 +386,10 @@ def __eq__(self, other): """Determines whether two TensorBase objects are equal using their associated keys. """ + if isinstance(other, (int, float)) and other == 0: + if isinstance(self, Tensor): + return isinstance(self.form, ZeroBaseForm) or self.form.empty() + return False return self._key == other._key def __ne__(self, other): @@ -452,13 +460,15 @@ def arg_function_spaces(self): """Returns a tuple of function spaces that the tensor is defined on. """ - return (self._function.ufl_function_space(),) + tensor = self._function + if isinstance(tensor, BaseForm): + return tuple(a.function_space() for a in tensor.arguments()) + else: + return (tensor.ufl_function_space(),) @cached_property def _argument(self): """Generates a 'test function' associated with this class.""" - from firedrake.ufl_expr import TestFunction - V, = self.arg_function_spaces return TestFunction(V) @@ -539,7 +549,6 @@ def arg_function_spaces(self): @cached_property def _argument(self): """Generates a tuple of 'test function' associated with this class.""" - from firedrake.ufl_expr import TestFunction return tuple(TestFunction(fs) for fs in self.arg_function_spaces) def arguments(self): @@ -650,7 +659,7 @@ def __init__(self, tensor, indices): """Constructor for the Block class.""" super(Block, self).__init__() self.operands = (tensor,) - self._blocks = dict(enumerate(indices)) + self._blocks = dict(enumerate(map(as_tuple, indices))) self._indices = indices @cached_property @@ -664,25 +673,10 @@ def _split_arguments(self): """Splits the function space and stores the component spaces determined by the indices. """ - from firedrake.functionspace import FunctionSpace, MixedFunctionSpace - from firedrake.ufl_expr import Argument - tensor, = self.operands - nargs = [] - for i, arg in enumerate(tensor.arguments()): - V = arg.function_space() - V_is = V.subfunctions - idx = as_tuple(self._blocks[i]) - if len(idx) == 1: - fidx, = idx - W = V_is[fidx] - W = FunctionSpace(W.mesh(), W.ufl_element()) - else: - W = MixedFunctionSpace([V_is[fidx] for fidx in idx]) - - nargs.append(Argument(W, arg.number(), part=arg.part())) - - return tuple(nargs) + return tuple(type(a)(subspace(a.function_space(), self._blocks[i]), + a.number(), part=a.part()) + for i, a in enumerate(tensor.arguments())) @cached_property def arg_function_spaces(self): @@ -880,7 +874,7 @@ class Tensor(TensorBase): def __init__(self, form, diagonal=False): """Constructor for the Tensor class.""" - if not isinstance(form, Form): + if not isinstance(form, (Form, ZeroBaseForm)): if isinstance(form, Function): raise TypeError("Use AssembledVector instead of Tensor.") raise TypeError("Only UFL forms are acceptable inputs.") @@ -936,6 +930,10 @@ def subdomain_data(self): """ return self.form.subdomain_data() + def empty(self): + """Returns whether the form associated with the tensor is empty.""" + return self.form.empty() + def _output_string(self, prec=None): """Creates a string representation of the tensor.""" return ["S", "V", "M"][self.rank] + "_%d" % self.id @@ -1103,6 +1101,13 @@ def _output_string(self, prec=None): class Transpose(UnaryOp): """An abstract Slate class representing the transpose of a tensor.""" + def __new__(cls, A): + if A == 0: + return Tensor(ZeroBaseForm(A.arguments()[::-1])) + if isinstance(A, Transpose): + tensor, = A.operands + return tensor + return BinaryOp.__new__(cls) @cached_property def arg_function_spaces(self): @@ -1127,6 +1132,10 @@ def _output_string(self, prec=None): class Negative(UnaryOp): """Abstract Slate class representing the negation of a tensor object.""" + def __new__(cls, A): + if A == 0: + return A + return BinaryOp.__new__(cls) @cached_property def arg_function_spaces(self): @@ -1197,6 +1206,12 @@ class Add(BinaryOp): :arg A: a :class:`~.firedrake.slate.TensorBase` object. :arg B: another :class:`~.firedrake.slate.TensorBase` object. """ + def __new__(cls, A, B): + if A == 0: + return B + elif B == 0: + return A + return BinaryOp.__new__(cls) def __init__(self, A, B): """Constructor for the Add class.""" @@ -1204,8 +1219,8 @@ def __init__(self, A, B): raise ValueError("Illegal op on a %s-tensor with a %s-tensor." % (A.shape, B.shape)) - assert all([space_equivalence(fsA, fsB) for fsA, fsB in - zip(A.arg_function_spaces, B.arg_function_spaces)]), ( + assert all(space_equivalence(fsA, fsB) for fsA, fsB in + zip(A.arg_function_spaces, B.arg_function_spaces)), ( "Function spaces associated with operands must match." ) @@ -1238,6 +1253,10 @@ class Mul(BinaryOp): :arg A: a :class:`~.firedrake.slate.TensorBase` object. :arg B: another :class:`~.firedrake.slate.TensorBase` object. """ + def __new__(cls, A, B): + if A == 0 or B == 0: + return Tensor(ZeroBaseForm(A.arguments()[:-1] + B.arguments()[1:])) + return BinaryOp.__new__(cls) def __init__(self, A, B): """Constructor for the Mul class.""" @@ -1288,14 +1307,14 @@ class Solve(BinaryOp): def __new__(cls, A, B, decomposition=None): assert A.rank == 2, "Operator must be a matrix." + assert B.rank >= 1, "RHS must be a vector or matrix." # Same rules for performing multiplication on Slate tensors # applies here. if A.shape[1] != B.shape[0]: - raise ValueError("Illegal op on a %s-tensor with a %s-tensor." - % (A.shape, B.shape)) + raise ValueError(f"Illegal op on a {A.shape}-tensor with a {B.shape}-tensor.") - fsA = A.arg_function_spaces[::-1][-1] + fsA = A.arg_function_spaces[0] fsB = B.arg_function_spaces[0] assert space_equivalence(fsA, fsB), ( @@ -1348,6 +1367,11 @@ class DiagonalTensor(UnaryOp): """ diagonal = True + def __new__(cls, A): + if A == 0: + return Tensor(ZeroBaseForm(A.arguments()[:1])) + return BinaryOp.__new__(cls) + def __init__(self, A): """Constructor for the Diagonal class.""" assert A.rank == 2, "The tensor must be rank 2." diff --git a/firedrake/slate/static_condensation/scpc.py b/firedrake/slate/static_condensation/scpc.py index 35fa4742eb..c4bda6dc27 100644 --- a/firedrake/slate/static_condensation/scpc.py +++ b/firedrake/slate/static_condensation/scpc.py @@ -102,7 +102,7 @@ def initialize(self, pc): r_expr = reduced_sys.rhs # Construct the condensed right-hand side - self._assemble_Srhs = get_assembler(r_expr, bcs=bcs, zero_bc_nodes=True, form_compiler_parameters=self.cxt.fc_params).assemble + self._assemble_Srhs = get_assembler(r_expr, bcs=bcs, form_compiler_parameters=self.cxt.fc_params).assemble # Allocate and set the condensed operator form_assembler = get_assembler(S_expr, bcs=bcs, form_compiler_parameters=self.cxt.fc_params, mat_type=mat_type, options_prefix=prefix, appctx=self.get_appctx(pc)) diff --git a/firedrake/solving_utils.py b/firedrake/solving_utils.py index 9e843016b5..8f03fbaf41 100644 --- a/firedrake/solving_utils.py +++ b/firedrake/solving_utils.py @@ -12,8 +12,8 @@ def _make_reasons(reasons): - return dict([(getattr(reasons, r), r) - for r in dir(reasons) if not r.startswith('_')]) + return {getattr(reasons, r): r + for r in dir(reasons) if not r.startswith('_')} KSPReasons = _make_reasons(PETSc.KSP.ConvergedReason()) @@ -221,7 +221,7 @@ def __init__(self, problem, mat_type, pmat_type, appctx=None, self._assemble_residual = get_assembler(self.F, bcs=self.bcs_F, form_compiler_parameters=self.fcp, - zero_bc_nodes=True).assemble + ).assemble self._jacobian_assembled = False self._splits = {} @@ -332,20 +332,17 @@ def split(self, fields): subu = function.Function(V, val=val) # Split it apart to shove in the form. subsplit = split(subu) - # Permutation from field indexing to indexing of pieces - field_renumbering = dict([f, i] for i, f in enumerate(field)) vec = [] for i, u in enumerate(us): if i in field: # If this is a field we're keeping, get it from # the new function. Otherwise just point to the # old data. - u = subsplit[field_renumbering[i]] + u = subsplit[field.index(i)] if u.ufl_shape == (): vec.append(u) else: - for idx in numpy.ndindex(u.ufl_shape): - vec.append(u[idx]) + vec.extend(u[idx] for idx in numpy.ndindex(u.ufl_shape)) # So now we have a new representation for the solution # vector in the old problem. For the fields we're going diff --git a/firedrake/tsfc_interface.py b/firedrake/tsfc_interface.py index ba10d79507..1117f54bd4 100644 --- a/firedrake/tsfc_interface.py +++ b/firedrake/tsfc_interface.py @@ -11,7 +11,7 @@ import ufl import finat.ufl -from ufl import Form, conj +from ufl import conj, Form, ZeroBaseForm from .ufl_expr import TestFunction from tsfc import compile_form as original_tsfc_compile_form @@ -203,7 +203,7 @@ def compile_form(form, name, parameters=None, split=True, interface=None, diagon iterable = ([(None, )*nargs, form], ) for idx, f in iterable: f = _real_mangle(f) - if not f.integrals(): + if isinstance(f, ZeroBaseForm) or f.empty(): # If we're assembling the R space component of a mixed argument, # and that component doesn't actually appear in the form then we # have an empty form, which we should not attempt to assemble. diff --git a/firedrake/variational_solver.py b/firedrake/variational_solver.py index 4a1ac396c5..3c8fc8b930 100644 --- a/firedrake/variational_solver.py +++ b/firedrake/variational_solver.py @@ -10,10 +10,11 @@ DEFAULT_SNES_PARAMETERS ) from firedrake.function import Function -from firedrake.ufl_expr import TrialFunction, TestFunction +from firedrake.ufl_expr import TrialFunction, TestFunction, action from firedrake.bcs import DirichletBC, EquationBC, extract_subdomain_ids, restricted_function_space from firedrake.adjoint_utils import NonlinearVariationalProblemMixin, NonlinearVariationalSolverMixin -from ufl import replace +from firedrake.__future__ import interpolate +from ufl import replace, Form __all__ = ["LinearVariationalProblem", "LinearVariationalSolver", @@ -91,8 +92,11 @@ def __init__(self, F, u, bcs=None, J=None, bcs = [bc.reconstruct(V=V_res, indices=bc._indices) for bc in bcs] self.u_restrict = Function(V_res).interpolate(u) v_res, u_res = TestFunction(V_res), TrialFunction(V_res) - F_arg, = F.arguments() - self.F = replace(F, {F_arg: v_res, self.u: self.u_restrict}) + if isinstance(F, Form): + F_arg, = F.arguments() + self.F = replace(F, {F_arg: v_res, self.u: self.u_restrict}) + else: + self.F = action(replace(F, {self.u: self.u_restrict}), interpolate(v_res, V)) v_arg, u_arg = self.J.arguments() self.J = replace(self.J, {v_arg: v_res, u_arg: u_res, self.u: self.u_restrict}) if self.Jp: diff --git a/tests/firedrake/equation_bcs/test_equation_bcs.py b/tests/firedrake/equation_bcs/test_equation_bcs.py index 087b07aa36..fdd05b7f2e 100644 --- a/tests/firedrake/equation_bcs/test_equation_bcs.py +++ b/tests/firedrake/equation_bcs/test_equation_bcs.py @@ -17,15 +17,13 @@ def nonlinear_poisson(solver_parameters, mesh_num, porder): u = Function(V) v = TestFunction(V) - f = Function(V) x, y = SpatialCoordinate(mesh) - f.interpolate(- 8.0 * pi * pi * cos(x * pi * 2) * cos(y * pi * 2)) + f = - 8.0 * pi * pi * cos(x * pi * 2) * cos(y * pi * 2) a = - inner(grad(u), grad(v)) * dx L = inner(f, v) * dx - g = Function(V) - g.interpolate(cos(2 * pi * x) * cos(2 * pi * y)) + g = cos(2 * pi * x) * cos(2 * pi * y) # Equivalent to bc1 = EquationBC(v * (u - g1) * ds(1) == 0, u, 1) e2 = as_vector([0., 1.]) @@ -33,7 +31,7 @@ def nonlinear_poisson(solver_parameters, mesh_num, porder): solve(a - L == 0, u, bcs=[bc1], solver_parameters=solver_parameters) - f.interpolate(cos(x * pi * 2) * cos(y * pi * 2)) + f = cos(x * pi * 2) * cos(y * pi * 2) return sqrt(assemble(inner(u - f, u - f) * dx)) @@ -46,15 +44,13 @@ def linear_poisson(solver_parameters, mesh_num, porder): u = TrialFunction(V) v = TestFunction(V) - f = Function(V) x, y = SpatialCoordinate(mesh) - f.interpolate(- 8.0 * pi * pi * cos(x * pi * 2) * cos(y * pi * 2)) + f = - 8.0 * pi * pi * cos(x * pi * 2) * cos(y * pi * 2) a = - inner(grad(u), grad(v)) * dx L = inner(f, v) * dx - g = Function(V) - g.interpolate(cos(2 * pi * x) * cos(2 * pi * y)) + g = cos(2 * pi * x) * cos(2 * pi * y) u_ = Function(V) @@ -62,7 +58,7 @@ def linear_poisson(solver_parameters, mesh_num, porder): solve(a == L, u_, bcs=[bc1], solver_parameters=solver_parameters) - f.interpolate(cos(x * pi * 2) * cos(y * pi * 2)) + f = cos(x * pi * 2) * cos(y * pi * 2) return sqrt(assemble(inner(u_ - f, u_ - f) * dx)) @@ -75,9 +71,8 @@ def nonlinear_poisson_bbc(solver_parameters, mesh_num, porder): u = Function(V) v = TestFunction(V) - f = Function(V) x, y = SpatialCoordinate(mesh) - f.interpolate(- 8.0 * pi * pi * cos(x * pi * 2)*cos(y * pi * 2)) + f = - 8.0 * pi * pi * cos(x * pi * 2)*cos(y * pi * 2) a = - inner(grad(u), grad(v)) * dx L = inner(f, v) * dx @@ -85,13 +80,13 @@ def nonlinear_poisson_bbc(solver_parameters, mesh_num, porder): e2 = as_vector([0., 1.]) a1 = (-inner(dot(grad(u), e2), dot(grad(v), e2)) + 4 * pi * pi * inner(u, v)) * ds(1) - g = Function(V).interpolate(cos(2 * pi * x) * cos(2 * pi * y)) + g = cos(2 * pi * x) * cos(2 * pi * y) bbc = DirichletBC(V, g, ((1, 3), (1, 4))) bc1 = EquationBC(a1 == 0, u, 1, bcs=[bbc]) solve(a - L == 0, u, bcs=[bc1], solver_parameters=solver_parameters) - f.interpolate(cos(x * pi * 2) * cos(y * pi * 2)) + f = cos(x * pi * 2) * cos(y * pi * 2) return sqrt(assemble(inner(u - f, u - f) * dx)) @@ -104,9 +99,8 @@ def linear_poisson_bbc(solver_parameters, mesh_num, porder): u = TrialFunction(V) v = TestFunction(V) - f = Function(V) x, y = SpatialCoordinate(mesh) - f.interpolate(- 8.0 * pi * pi * cos(x * pi * 2)*cos(y * pi * 2)) + f = - 8.0 * pi * pi * cos(x * pi * 2)*cos(y * pi * 2) a = - inner(grad(u), grad(v)) * dx L = inner(f, v) * dx @@ -117,13 +111,13 @@ def linear_poisson_bbc(solver_parameters, mesh_num, porder): u = Function(V) - g = Function(V).interpolate(cos(2 * pi * x) * cos(2 * pi * y)) + g = cos(2 * pi * x) * cos(2 * pi * y) bbc = DirichletBC(V, g, ((1, 3), (1, 4))) bc1 = EquationBC(a1 == L1, u, 1, bcs=[bbc]) solve(a == L, u, bcs=[bc1], solver_parameters=solver_parameters) - f.interpolate(cos(x * pi * 2)*cos(y * pi * 2)) + f = cos(x * pi * 2)*cos(y * pi * 2) return sqrt(assemble(inner(u - f, u - f) * dx)) @@ -141,22 +135,25 @@ def nonlinear_poisson_mixed(solver_parameters, mesh_num, porder): n = FacetNormal(mesh) x, y = SpatialCoordinate(mesh) - f = Function(DG).interpolate(-8 * pi * pi * cos(2 * pi * x + pi / 3) * cos(2 * pi * y)) - u1 = Function(DG).interpolate(cos(2 * pi * y) / 2) + f = -8 * pi * pi * cos(2 * pi * x + pi / 3) * cos(2 * pi * y) + u1 = cos(2 * pi * y) / 2 - a = (inner(sigma, tau) + inner(u, div(tau)) + inner(div(sigma), v)) * dx + a = inner(sigma, tau) * dx + inner(u, div(tau)) * dx + inner(div(sigma), v) * dx L = inner(u1, dot(tau, n)) * ds(1) + inner(f, v) * dx - g = Function(BDM).project(as_vector([-2 * pi * sin(2 * pi * x + pi / 3) * cos(2 * pi * y), -2 * pi * cos(2 * pi * x + pi / 3) * sin(2 * pi * y)])) + g = as_vector([-2 * pi * sin(2 * pi * x + pi / 3) * cos(2 * pi * y), -2 * pi * cos(2 * pi * x + pi / 3) * sin(2 * pi * y)]) - bc2 = EquationBC(inner((dot(sigma, n) - dot(g, n)), dot(tau, n)) * ds(2) == 0, w, 2, V=W.sub(0)) - bc3 = EquationBC(inner((dot(sigma, n) - dot(g, n)), dot(tau, n)) * ds(3) == 0, w, 3, V=W.sub(0)) + tau_n = dot(tau, n) + sig_n = dot(sigma, n) + g_n = dot(g, n) + bc2 = EquationBC(inner(sig_n - g_n, tau_n) * ds(2) == 0, w, 2, V=W.sub(0)) + bc3 = EquationBC(inner(sig_n - g_n, tau_n) * ds(3) == 0, w, 3, V=W.sub(0)) bc4 = DirichletBC(W.sub(0), g, 4) solve(a - L == 0, w, bcs=[bc2, bc3, bc4], solver_parameters=solver_parameters) - f.interpolate(cos(2 * pi * x + pi / 3) * cos(2 * pi * y)) - g = Function(BDM).project(as_vector([-2 * pi * sin(2 * pi * x + pi / 3) * cos(2 * pi * y), -2 * pi * cos(2 * pi * x + pi / 3) * sin(2 * pi * y)])) + f = cos(2 * pi * x + pi / 3) * cos(2 * pi * y) + g = as_vector([-2 * pi * sin(2 * pi * x + pi / 3) * cos(2 * pi * y), -2 * pi * cos(2 * pi * x + pi / 3) * sin(2 * pi * y)]) return sqrt(assemble(inner(u - f, u - f) * dx)), sqrt(assemble(inner(sigma - g, sigma - g) * dx)) @@ -173,28 +170,30 @@ def linear_poisson_mixed(solver_parameters, mesh_num, porder): tau, v = TestFunctions(W) x, y = SpatialCoordinate(mesh) - f = Function(DG).interpolate(-8 * pi * pi * cos(2 * pi * x + pi / 3) * cos(2 * pi * y)) - u1 = Function(DG).interpolate(cos(2 * pi * y) / 2) + f = -8 * pi * pi * cos(2 * pi * x + pi / 3) * cos(2 * pi * y) + u1 = cos(2 * pi * y) / 2 n = FacetNormal(mesh) a = (inner(sigma, tau) + inner(u, div(tau)) + inner(div(sigma), v)) * dx L = inner(u1, dot(tau, n)) * ds(1) + inner(f, v) * dx - g = Function(BDM).project(as_vector([-2 * pi * sin(2 * pi * x + pi / 3) * cos(2 * pi * y), -2 * pi * cos(2 * pi * x + pi / 3) * sin(2 * pi * y)])) + g = as_vector([-2 * pi * sin(2 * pi * x + pi / 3) * cos(2 * pi * y), -2 * pi * cos(2 * pi * x + pi / 3) * sin(2 * pi * y)]) w = Function(W) - bc2 = EquationBC(inner(n, tau) * inner(sigma, n) * ds(2) == inner(n, tau) * inner(g, n) * ds(2), w, 2, V=W.sub(0)) - bc3 = EquationBC(inner(n, tau) * inner(sigma, n) * ds(3) == inner(n, tau) * inner(g, n) * ds(3), w, 3, V=W.sub(0)) + tau_n = dot(tau, n) + sig_n = dot(sigma, n) + g_n = dot(g, n) + bc2 = EquationBC(inner(sig_n, tau_n) * ds(2) == inner(g_n, tau_n) * ds(2), w, 2, V=W.sub(0)) + bc3 = EquationBC(inner(sig_n, tau_n) * ds(3) == inner(g_n, tau_n) * ds(3), w, 3, V=W.sub(0)) bc4 = DirichletBC(W.sub(0), g, 4) solve(a == L, w, bcs=[bc2, bc3, bc4], solver_parameters=solver_parameters) - sigma, u = w.subfunctions - - f.interpolate(cos(2 * pi * x + pi / 3) * cos(2 * pi * y)) - g = Function(BDM).project(as_vector([-2 * pi * sin(2 * pi * x + pi / 3) * cos(2 * pi * y), -2 * pi * cos(2 * pi * x + pi / 3) * sin(2 * pi * y)])) + f = cos(2 * pi * x + pi / 3) * cos(2 * pi * y) + g = as_vector([-2 * pi * sin(2 * pi * x + pi / 3) * cos(2 * pi * y), -2 * pi * cos(2 * pi * x + pi / 3) * sin(2 * pi * y)]) + sigma, u = w.subfunctions return sqrt(assemble(inner(u - f, u - f) * dx)), sqrt(assemble(inner(sigma - g, sigma - g) * dx)) @@ -202,7 +201,7 @@ def linear_poisson_mixed(solver_parameters, mesh_num, porder): @pytest.mark.parametrize("with_bbc", [False, True]) def test_EquationBC_poisson_matrix(eq_type, with_bbc): mat_type = "aij" - porder = 3 + porder = 2 # Test standard poisson with EquationBCs # aij @@ -235,7 +234,7 @@ def test_EquationBC_poisson_matrix(eq_type, with_bbc): def test_EquationBC_poisson_matfree(with_bbc): eq_type = "linear" mat_type = "matfree" - porder = 3 + porder = 2 # Test standard poisson with EquationBCs # matfree @@ -271,7 +270,7 @@ def test_EquationBC_poisson_matfree(with_bbc): @pytest.mark.parametrize("eq_type", ["linear", "nonlinear"]) def test_EquationBC_mixedpoisson_matrix(eq_type): mat_type = "aij" - porder = 2 + porder = 0 # Mixed poisson with EquationBCs # aij @@ -294,7 +293,7 @@ def test_EquationBC_mixedpoisson_matrix(eq_type): def test_EquationBC_mixedpoisson_matrix_fieldsplit(): mat_type = "aij" eq_type = "linear" - porder = 2 + porder = 0 # Mixed poisson with EquationBCs # aij with fieldsplit pc @@ -324,7 +323,7 @@ def test_EquationBC_mixedpoisson_matrix_fieldsplit(): def test_EquationBC_mixedpoisson_matfree_fieldsplit(): mat_type = "matfree" eq_type = "linear" - porder = 2 + porder = 0 # Mixed poisson with EquationBCs # matfree with fieldsplit pc @@ -366,7 +365,7 @@ def test_equation_bcs_pc(): v, w = split(TestFunction(V)) x, y = SpatialCoordinate(mesh) exact = cos(2 * pi * x) * cos(2 * pi * y) - g = Function(CG).interpolate(8 * pi**2 * exact) + g = 8 * pi**2 * exact F = inner(grad(u), grad(v)) * dx + inner(l, w) * dx - inner(g, v) * dx bc = EquationBC(inner((u - exact), v) * ds == 0, f, (1, 2, 3, 4), V=V.sub(0)) params = { diff --git a/tests/firedrake/multigrid/test_poisson_gmg.py b/tests/firedrake/multigrid/test_poisson_gmg.py index 81f56acbfc..577587a393 100644 --- a/tests/firedrake/multigrid/test_poisson_gmg.py +++ b/tests/firedrake/multigrid/test_poisson_gmg.py @@ -195,12 +195,11 @@ def test_baseform_coarsening(solver_type, mixed): a_terms.append(inner(grad(u), grad(v)) * dx) a = sum(a_terms) - assemble_bcs = lambda L: assemble(L, bcs=bcs, zero_bc_nodes=True) # These are equivalent right-hand sides sources = [sum(forms), # purely symbolic linear form - assemble_bcs(sum(forms)), # purely numerical cofunction - sum(assemble_bcs(form) for form in forms), # symbolic combination of numerical cofunctions - forms[0] + assemble_bcs(sum(forms[1:])), # symbolic plus numerical + assemble(sum(forms), bcs=bcs), # purely numerical cofunction + sum(assemble(form, bcs=bcs) for form in forms), # symbolic combination of numerical cofunctions + forms[0] + assemble(sum(forms[1:]), bcs=bcs), # symbolic plus numerical ] solutions = [] for L in sources: diff --git a/tests/firedrake/regression/test_assemble.py b/tests/firedrake/regression/test_assemble.py index bd8f020e60..72394380cb 100644 --- a/tests/firedrake/regression/test_assemble.py +++ b/tests/firedrake/regression/test_assemble.py @@ -225,7 +225,7 @@ def test_one_form_assembler_cache(mesh): assert len(L._cache[_FORM_CACHE_KEY]) == 3 # changing zero_bc_nodes should increase the cache size - assemble(L, zero_bc_nodes=True) + assemble(L, zero_bc_nodes=False) assert len(L._cache[_FORM_CACHE_KEY]) == 4 diff --git a/tests/firedrake/regression/test_assemble_baseform.py b/tests/firedrake/regression/test_assemble_baseform.py index 063a33bdd9..54cc0b183c 100644 --- a/tests/firedrake/regression/test_assemble_baseform.py +++ b/tests/firedrake/regression/test_assemble_baseform.py @@ -1,7 +1,7 @@ import pytest import numpy as np from firedrake import * -from firedrake.assemble import BaseFormAssembler, get_assembler +from firedrake.assemble import get_assembler from firedrake.utils import ScalarType import ufl @@ -43,39 +43,15 @@ def fs(request, mesh): @pytest.fixture def f(fs): f = Function(fs, name="f") - f_split = f.subfunctions x = SpatialCoordinate(fs.mesh())[0] - - # NOTE: interpolation of UFL expressions into mixed - # function spaces is not yet implemented - for fi in f_split: - fs_i = fi.function_space() - if fs_i.rank == 1: - fi.interpolate(as_vector((x,) * fs_i.value_size)) - elif fs_i.rank == 2: - fi.interpolate(as_tensor([[x for i in range(fs_i.mesh().geometric_dimension())] - for j in range(fs_i.rank)])) - else: - fi.interpolate(x) + f.interpolate(as_tensor(np.full(f.ufl_shape, x))) return f @pytest.fixture def one(fs): one = Function(fs, name="one") - ones = one.subfunctions - - # NOTE: interpolation of UFL expressions into mixed - # function spaces is not yet implemented - for fi in ones: - fs_i = fi.function_space() - if fs_i.rank == 1: - fi.interpolate(Constant((1.0,) * fs_i.value_size)) - elif fs_i.rank == 2: - fi.interpolate(Constant([[1.0 for i in range(fs_i.mesh().geometric_dimension())] - for j in range(fs_i.rank)])) - else: - fi.interpolate(Constant(1.0)) + one.interpolate(Constant(np.ones(one.ufl_shape))) return one @@ -155,23 +131,6 @@ def test_zero_form(M, f, one): assert abs(zero_form - 0.5 * np.prod(f.ufl_shape)) < 1.0e-12 -def test_preprocess_form(M, a, f): - from ufl.algorithms import expand_indices, expand_derivatives - - expr = action(action(M, M), f) - A = BaseFormAssembler.preprocess_base_form(expr) - B = action(expand_derivatives(M), action(M, f)) - - assert isinstance(A, ufl.Action) - try: - # Need to expand indices to be able to match equal (different MultiIndex used for both). - assert expand_indices(A.left()) == expand_indices(B.left()) - assert expand_indices(A.right()) == expand_indices(B.right()) - except KeyError: - # Index expansion doesn't seem to play well with tensor elements. - pass - - def test_tensor_copy(a, M): # 1-form tensor diff --git a/tests/firedrake/regression/test_bcs.py b/tests/firedrake/regression/test_bcs.py index 0e27515890..16cbc669c6 100644 --- a/tests/firedrake/regression/test_bcs.py +++ b/tests/firedrake/regression/test_bcs.py @@ -327,7 +327,7 @@ def test_bcs_rhs_assemble(a, V): b1 = assemble(a) b1_func = b1.riesz_representation(riesz_map="l2") for bc in bcs: - bc.apply(b1_func) + bc.zero(b1_func) b1.assign(b1_func.riesz_representation(riesz_map="l2")) b2 = assemble(a, bcs=bcs) assert np.allclose(b1.dat.data, b2.dat.data) diff --git a/tests/firedrake/regression/test_netgen.py b/tests/firedrake/regression/test_netgen.py index 904e9a6819..8cc164d98c 100644 --- a/tests/firedrake/regression/test_netgen.py +++ b/tests/firedrake/regression/test_netgen.py @@ -51,8 +51,7 @@ def poisson(h, degree=2): # Assembling matrix A = assemble(a, bcs=bc) - b = assemble(l) - bc.apply(b) + b = assemble(l, bcs=bc) # Solving the problem solve(A, u, b, solver_parameters={"ksp_type": "preonly", "pc_type": "lu"}) @@ -95,8 +94,7 @@ def poisson3D(h, degree=2): # Assembling matrix A = assemble(a, bcs=bc) - b = assemble(l) - bc.apply(b) + b = assemble(l, bcs=bc) # Solving the problem solve(A, u, b, solver_parameters={"ksp_type": "preonly", "pc_type": "lu"}) diff --git a/tests/firedrake/regression/test_restricted_function_space.py b/tests/firedrake/regression/test_restricted_function_space.py index dc9a2ecc64..6dd093f58a 100644 --- a/tests/firedrake/regression/test_restricted_function_space.py +++ b/tests/firedrake/regression/test_restricted_function_space.py @@ -146,7 +146,8 @@ def test_poisson_inhomogeneous_bcs_2(j): @pytest.mark.parallel(nprocs=3) -def test_poisson_inhomogeneous_bcs_high_level_interface(): +@pytest.mark.parametrize("assembled_rhs", [False, True], ids=("Form", "Cofunction")) +def test_poisson_inhomogeneous_bcs_high_level_interface(assembled_rhs): mesh = UnitSquareMesh(8, 8) V = FunctionSpace(mesh, "CG", 2) bc1 = DirichletBC(V, 0., 1) @@ -155,9 +156,11 @@ def test_poisson_inhomogeneous_bcs_high_level_interface(): v = TestFunction(V) a = inner(grad(u), grad(v)) * dx u = Function(V) - L = inner(Constant(0), v) * dx + L = inner(Constant(-2), v) * dx + if assembled_rhs: + L = assemble(L) solve(a == L, u, bcs=[bc1, bc2], restrict=True) - assert errornorm(SpatialCoordinate(mesh)[0], u) < 1.e-12 + assert errornorm(SpatialCoordinate(mesh)[0]**2, u) < 1.e-12 @pytest.mark.parametrize("j", [1, 2, 5]) diff --git a/tests/firedrake/regression/test_solving_interface.py b/tests/firedrake/regression/test_solving_interface.py index f32e6f3214..ed47d9e82e 100644 --- a/tests/firedrake/regression/test_solving_interface.py +++ b/tests/firedrake/regression/test_solving_interface.py @@ -222,22 +222,24 @@ def test_constant_jacobian_lvs(): def test_solve_cofunction_rhs(): - mesh = UnitSquareMesh(10, 10) + mesh = UnitIntervalMesh(10) V = FunctionSpace(mesh, "CG", 1) + x, = SpatialCoordinate(mesh) u = TrialFunction(V) v = TestFunction(V) - a = inner(u, v) * dx - + a = inner(grad(u), grad(v)) * dx L = Cofunction(V.dual()) - L.vector()[:] = 1. + bcs = [DirichletBC(V, x, "on_boundary")] + # Set the wrong BCs on the RHS + for bc in bcs: + bc.set(L, 888) + Lold = L.copy() w = Function(V) - solve(a == L, w) - - Aw = assemble(action(a, w)) - assert isinstance(Aw, Cofunction) - assert np.allclose(Aw.dat.data_ro, L.dat.data_ro) + solve(a == L, w, bcs=bcs) + assert errornorm(x, w) < 1E-10 + assert np.allclose(L.dat.data, Lold.dat.data) def test_solve_empty_form_rhs(): diff --git a/tests/firedrake/slate/test_assemble_tensors.py b/tests/firedrake/slate/test_assemble_tensors.py index 5aff159b9b..c35d43e27e 100644 --- a/tests/firedrake/slate/test_assemble_tensors.py +++ b/tests/firedrake/slate/test_assemble_tensors.py @@ -249,9 +249,13 @@ def test_matrix_subblocks(mesh): refs = dict(split_form(A.form)) _A = A.blocks for x, y in indices: - ref = assemble(refs[x, y]).M.values block = _A[x, y] - assert np.allclose(assemble(block).M.values, ref, rtol=1e-14) + ref = refs[x, y] + if isinstance(ref, Form): + assert np.allclose(assemble(block).M.values, + assemble(ref).M.values, rtol=1e-14) + elif isinstance(ref, ZeroBaseForm): + assert block.form == ref # Mixed blocks A0101 = _A[:2, :2] @@ -280,9 +284,12 @@ def test_matrix_subblocks(mesh): (A1212_10, refs[(2, 1)])] # Test assembly of blocks of mixed blocks - for tensor, form in items: - ref = assemble(form).M.values - assert np.allclose(assemble(tensor).M.values, ref, rtol=1e-14) + for block, ref in items: + if isinstance(ref, Form): + assert np.allclose(assemble(block).M.values, + assemble(ref).M.values, rtol=1e-14) + elif isinstance(ref, ZeroBaseForm): + assert block.form == ref def test_diagonal(mass, matrix_mixed_nofacet):