diff --git a/firedrake/formmanipulation.py b/firedrake/formmanipulation.py index 3179961df8..708adfd8e5 100644 --- a/firedrake/formmanipulation.py +++ b/firedrake/formmanipulation.py @@ -2,23 +2,22 @@ import numpy import collections -from ufl import as_vector, split, ZeroBaseForm -from ufl.classes import Zero, FixedIndex, ListTensor +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, FunctionSpace +from firedrake.cofunction import Cofunction +from firedrake.functionspace import FunctionSpace, MixedFunctionSpace, DualSpace def subspace(V, indices): - try: - indices = tuple(indices) - except TypeError: - # Only one index provided. - indices = (indices, ) if len(indices) == 1: W = V[indices[0]] W = FunctionSpace(W.mesh(), W.ufl_element()) @@ -66,7 +65,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 @@ -75,11 +74,13 @@ def split(self, form, argument_indices): assert (idx[0] == 0 for idx in self.blocks.values()) return form f = map_integrand_dags(self, form) - f = expand_derivatives(f) - if f.empty(): - f = ZeroBaseForm(tuple(Argument(subspace(arg.function_space(), indices), + # TODO find a way to distinguish empty Forms avoiding expand_derivatives + if expand_derivatives(f).empty(): + # Get ZeroBaseForm with the right shape + f = ZeroBaseForm(tuple(Argument(subspace(arg.function_space(), + self.blocks[arg.number()]), arg.number(), part=arg.part()) - for arg, indices in zip(form.arguments(), argument_indices))) + for arg in form.arguments())) return f expr = MultiFunction.reuse_if_untouched @@ -109,6 +110,7 @@ def coefficient_derivative(self, o, expr, coefficients, arguments, cds): @PETSc.Log.EventDecorator() def argument(self, o): V = o.function_space() + if len(V) == 1: # Not on a mixed space, just return ourselves. return o @@ -118,12 +120,6 @@ def argument(self, o): indices = self.blocks[o.number()] - try: - indices = tuple(indices) - except TypeError: - # Only one index provided. - indices = (indices, ) - W = subspace(V, indices) a = Argument(W, o.number(), part=o.part()) a = (a, ) if len(W) == 1 else split(a) @@ -141,6 +137,28 @@ def argument(self, o): 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() + + if len(V) == 1: + # Not on a mixed space, just return ourselves. + return o + + try: + indices, = set(self.blocks.values()) + except ValueError: + raise ValueError("Cofunction found on an off-diagonal block") + + if len(indices) == 1: + i = indices[0] + W = V[i] + W = DualSpace(W.mesh(), W.ufl_element()) + c = Cofunction(W, val=o.dat[i]) + else: + W = MixedFunctionSpace([V[i] for i in indices]) + c = Cofunction(W, val=MixedDat(o.dat[i] for i in indices)) + return c + SplitForm = collections.namedtuple("SplitForm", ["indices", "form"]) diff --git a/scripts/firedrake-install b/scripts/firedrake-install index fa39e2b3b0..17a8d7fdbc 100755 --- a/scripts/firedrake-install +++ b/scripts/firedrake-install @@ -16,10 +16,21 @@ from glob import iglob from itertools import chain import re import importlib + + +class InstallError(Exception): + # Exception for generic install problems. + pass + + try: from pkg_resources.extern.packaging.version import Version, InvalidVersion except ModuleNotFoundError: - from packaging.version import Version, InvalidVersion + try: + from packaging.version import Version, InvalidVersion + except ModuleNotFoundError: + raise InstallError("Neither setuptools or packaging found. Please " + "install one of these packages before trying again.") osname = platform.uname().system arch = platform.uname().machine @@ -52,11 +63,6 @@ firedrake_apps = { } -class InstallError(Exception): - # Exception for generic install problems. - pass - - class FiredrakeConfiguration(dict): """A dictionary extended to facilitate the storage of Firedrake configuration information.""" diff --git a/tests/firedrake/regression/test_linesmoother.py b/tests/firedrake/regression/test_linesmoother.py index ba1def7c33..049421ddcc 100644 --- a/tests/firedrake/regression/test_linesmoother.py +++ b/tests/firedrake/regression/test_linesmoother.py @@ -43,7 +43,8 @@ def backend(request): return request.param -def test_linesmoother(mesh, S1family, expected, backend): +@pytest.mark.parametrize("rhs", ["form_rhs", "cofunc_rhs"]) +def test_linesmoother(mesh, S1family, expected, backend, rhs): base_cell = mesh._base_mesh.ufl_cell() S2family = "DG" if base_cell.is_simplex() else "DQ" DGfamily = "DG" if mesh.ufl_cell().is_simplex() else "DQ" @@ -86,6 +87,10 @@ def test_linesmoother(mesh, S1family, expected, backend): f = exp(-rsq) L = inner(f, q)*dx(degree=2*(degree+1)) + if rhs == 'cofunc_rhs': + L = assemble(L) + elif rhs != 'form_rhs': + raise ValueError("Unknown right hand side type") w0 = Function(W) problem = LinearVariationalProblem(a, L, w0, bcs=bcs, aP=aP, form_compiler_parameters={"mode": "vanilla"}) diff --git a/tests/firedrake/regression/test_matrix_free.py b/tests/firedrake/regression/test_matrix_free.py index 6cc6b207f7..f8528a1ed5 100644 --- a/tests/firedrake/regression/test_matrix_free.py +++ b/tests/firedrake/regression/test_matrix_free.py @@ -130,6 +130,7 @@ def test_matrixfree_action(a, V, bcs): @pytest.mark.parametrize("preassembled", [False, True], ids=["variational", "preassembled"]) +@pytest.mark.parametrize("rhs", ["form_rhs", "cofunc_rhs"]) @pytest.mark.parametrize("parameters", [{"ksp_type": "preonly", "pc_type": "python", @@ -168,7 +169,7 @@ def test_matrixfree_action(a, V, bcs): "fieldsplit_1_fieldsplit_1_pc_type": "python", "fieldsplit_1_fieldsplit_1_pc_python_type": "firedrake.AssembledPC", "fieldsplit_1_fieldsplit_1_assembled_pc_type": "lu"}]) -def test_fieldsplitting(mesh, preassembled, parameters): +def test_fieldsplitting(mesh, preassembled, parameters, rhs): V = FunctionSpace(mesh, "CG", 1) P = FunctionSpace(mesh, "DG", 0) Q = VectorFunctionSpace(mesh, "DG", 1) @@ -185,6 +186,10 @@ def test_fieldsplitting(mesh, preassembled, parameters): a = inner(u, v)*dx L = inner(expect, v)*dx + if rhs == 'cofunc_rhs': + L = assemble(L) + elif rhs != 'form_rhs': + raise ValueError("Unknown right hand side type") f = Function(W) diff --git a/tests/firedrake/regression/test_nullspace.py b/tests/firedrake/regression/test_nullspace.py index bdef9fd6a3..e4f436ff8d 100644 --- a/tests/firedrake/regression/test_nullspace.py +++ b/tests/firedrake/regression/test_nullspace.py @@ -291,7 +291,8 @@ def test_nullspace_mixed_multiple_components(): @pytest.mark.parallel(nprocs=2) @pytest.mark.parametrize("aux_pc", [False, True], ids=["PC(mu)", "PC(DG0-mu)"]) -def test_near_nullspace_mixed(aux_pc): +@pytest.mark.parametrize("rhs", ["form_rhs", "cofunc_rhs"]) +def test_near_nullspace_mixed(aux_pc, rhs): # test nullspace and nearnullspace for a mixed Stokes system # this is tested on the SINKER case of May and Moresi https://doi.org/10.1016/j.pepi.2008.07.036 # fails in parallel if nullspace is copied to fieldsplit_1_Mp_ksp solve (see PR #3488) @@ -323,6 +324,10 @@ def test_near_nullspace_mixed(aux_pc): f = as_vector((0, -9.8*conditional(inside_box, 2, 1))) L = inner(f, v)*dx + if rhs == 'cofunc_rhs': + L = assemble(L) + elif rhs != 'form_rhs': + raise ValueError("Unknown right hand side type") bcs = [DirichletBC(W[0].sub(0), 0, (1, 2)), DirichletBC(W[0].sub(1), 0, (3, 4))]