From 109b39dad95a507b726f8fc96c0f07f98792d9f9 Mon Sep 17 00:00:00 2001 From: David Ham Date: Tue, 17 Dec 2024 14:56:46 +0000 Subject: [PATCH 01/23] test for correct annotation of cofunction assign --- tests/firedrake/regression/test_adjoint_operators.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/firedrake/regression/test_adjoint_operators.py b/tests/firedrake/regression/test_adjoint_operators.py index 8d4c0ab796..d0abec9a0b 100644 --- a/tests/firedrake/regression/test_adjoint_operators.py +++ b/tests/firedrake/regression/test_adjoint_operators.py @@ -1002,7 +1002,6 @@ def test_cofunction_assign_functional(): f.assign(2.0) assert np.allclose(Jhat(f), 2.0) - @pytest.mark.skipcomplex # Taping for complex-valued 0-forms not yet done def test_bdy_control(): # Test for the case the boundary condition is a control for a From aa010bfde6759ada4d0f5b44c8164175878c4870 Mon Sep 17 00:00:00 2001 From: David Ham Date: Mon, 8 Jul 2024 16:06:06 +0100 Subject: [PATCH 02/23] DROP BEFORE MERGE --- .github/workflows/build.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 909db6990a..acebd51af8 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -83,6 +83,7 @@ jobs: --install defcon \ --install gadopt \ --install asQ \ + --package-branch ufl dham/cofunction_is_terminal \ || (cat firedrake-install.log && /bin/false) - name: Install test dependencies run: | From e56f5ffc3b5b78314d7900c020f2e2316b96f19d Mon Sep 17 00:00:00 2001 From: David Ham Date: Mon, 1 Jul 2024 13:27:24 +0100 Subject: [PATCH 03/23] beginnings of doing inner product right --- firedrake/adjoint_utils/function.py | 11 ++++++++++- firedrake/cofunction.py | 16 ++++++++-------- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 5e87751d36..af44b05990 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -316,7 +316,10 @@ def _ad_add(self, other): return r def _ad_dot(self, other, options=None): - from firedrake import assemble + from firedrake import assemble, action, Cofunction + + if isinstance(other, Cofunction): + return assemble(action(other, self)) options = {} if options is None else options riesz_representation = options.get("riesz_representation", "L2") @@ -406,3 +409,9 @@ def _ad_to_petsc(self, vec=None): def __deepcopy__(self, memodict={}): return self.copy(deepcopy=True) + + +class CofunctionMixin(FunctionMixin): + + def _ad_dot(self, other): + return assemble(firedrake.action(self, other)) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index 4878e6da59..bca08c8213 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -8,13 +8,13 @@ import firedrake.functionspaceimpl as functionspaceimpl from firedrake import utils, vector, ufl_expr from firedrake.utils import ScalarType -from firedrake.adjoint_utils.function import FunctionMixin +from firedrake.adjoint_utils.function import CofunctionMixin from firedrake.adjoint_utils.checkpointing import DelegatedFunctionCheckpoint from firedrake.adjoint_utils.blocks.function import CofunctionAssignBlock from firedrake.petsc import PETSc -class Cofunction(ufl.Cofunction, FunctionMixin): +class Cofunction(ufl.Cofunction, CofunctionMixin): r"""A :class:`Cofunction` represents a function on a dual space. Like Functions, cofunctions are represented as sums of basis functions: @@ -33,7 +33,7 @@ class Cofunction(ufl.Cofunction, FunctionMixin): """ @PETSc.Log.EventDecorator() - @FunctionMixin._ad_annotate_init + @CofunctionMixin._ad_annotate_init def __init__(self, function_space, val=None, name=None, dtype=ScalarType, count=None): r""" @@ -105,13 +105,13 @@ def _analyze_form_arguments(self): self._coefficients = (self,) @utils.cached_property - @FunctionMixin._ad_annotate_subfunctions + @CofunctionMixin._ad_annotate_subfunctions def subfunctions(self): r"""Extract any sub :class:`Cofunction`\s defined on the component spaces of this this :class:`Cofunction`'s :class:`.FunctionSpace`.""" return tuple(type(self)(fs, dat) for fs, dat in zip(self.function_space(), self.dat)) - @FunctionMixin._ad_annotate_subfunctions + @CofunctionMixin._ad_annotate_subfunctions def split(self): import warnings warnings.warn("The .split() method is deprecated, please use the .subfunctions property instead", category=FutureWarning) @@ -260,7 +260,7 @@ def riesz_representation(self, riesz_map='L2', **solver_options): "riesz_representation": riesz_map, "solver_options": solver_options}) - @FunctionMixin._ad_annotate_iadd + @CofunctionMixin._ad_annotate_iadd @utils.known_pyop2_safe def __iadd__(self, expr): @@ -276,7 +276,7 @@ def __iadd__(self, expr): # Let Python hit `BaseForm.__add__` which relies on ufl.FormSum. return NotImplemented - @FunctionMixin._ad_annotate_isub + @CofunctionMixin._ad_annotate_isub @utils.known_pyop2_safe def __isub__(self, expr): @@ -293,7 +293,7 @@ def __isub__(self, expr): # Let Python hit `BaseForm.__sub__` which relies on ufl.FormSum. return NotImplemented - @FunctionMixin._ad_annotate_imul + @CofunctionMixin._ad_annotate_imul def __imul__(self, expr): if np.isscalar(expr): From 6fc8cdbc5798025876a33c8bd58d6bde795f2c8b Mon Sep 17 00:00:00 2001 From: "David A. Ham" Date: Mon, 1 Jul 2024 17:51:02 +0100 Subject: [PATCH 04/23] rieszmap --- firedrake/adjoint_utils/blocks/solving.py | 8 +- firedrake/adjoint_utils/function.py | 136 ++++++++++++++++------ firedrake/cofunction.py | 4 +- firedrake/ufl_expr.py | 6 + 4 files changed, 115 insertions(+), 39 deletions(-) diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index e4664665b0..2028bc2684 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -756,7 +756,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, c = block_variable.output c_rep = block_variable.saved_output - if isinstance(c, firedrake.Function): + if isinstance(c, (firedrake.Function, firedrake.Cofunction)): trial_function = firedrake.TrialFunction(c.function_space()) elif isinstance(c, firedrake.Constant): mesh = F_form.ufl_domain() @@ -793,7 +793,11 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, replace_map[self.func] = self.get_outputs()[0].saved_output dFdm = replace(dFdm, replace_map) - dFdm = dFdm * adj_sol + if isinstance(dFdm, firedrake.Argument): + # Corner case. Should be fixed more permanenty upstream in UFL. + dFdm = ufl.Action(dFdm, adj_sol) + else: + dFdm = dFdm * adj_sol dFdm = firedrake.assemble(dFdm, **self.assemble_kwargs) return dFdm diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index af44b05990..8ef0fa6696 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -1,4 +1,4 @@ -from functools import wraps +from functools import wraps, cached_property import ufl from ufl.domain import extract_unique_domain from pyadjoint.overloaded_type import create_overloaded_object, FloatingType @@ -226,7 +226,6 @@ def _ad_convert_riesz(self, value, options=None): options = {} if options is None else options riesz_representation = options.get("riesz_representation", "L2") solver_options = options.get("solver_options", {}) - V = options.get("function_space", self.function_space()) if value == 0.: # In adjoint-based differentiation, value == 0. arises only when # the functional is independent on the control variable. @@ -235,41 +234,11 @@ def _ad_convert_riesz(self, value, options=None): if not isinstance(value, (Cofunction, Function)): raise TypeError("Expected a Cofunction or a Function") - if riesz_representation == "l2": - return Function(V, val=value.dat) - - elif riesz_representation in ("L2", "H1"): - if not isinstance(value, Cofunction): - raise TypeError("Expected a Cofunction") - - ret = Function(V) - a = self._define_riesz_map_form(riesz_representation, V) - firedrake.solve(a == value, ret, **solver_options) - return ret - - elif callable(riesz_representation): + if callable(riesz_representation): return riesz_representation(value) - else: - raise ValueError( - "Unknown Riesz representation %s" % riesz_representation) - - def _define_riesz_map_form(self, riesz_representation, V): - from firedrake import TrialFunction, TestFunction - - u = TrialFunction(V) - v = TestFunction(V) - if riesz_representation == "L2": - a = firedrake.inner(u, v)*firedrake.dx - - elif riesz_representation == "H1": - a = firedrake.inner(u, v)*firedrake.dx \ - + firedrake.inner(firedrake.grad(u), firedrake.grad(v))*firedrake.dx - - else: - raise NotImplementedError( - "Unknown Riesz representation %s" % riesz_representation) - return a + return RieszMap(self, riesz_representation, + solver_options=solver_options)(value) @no_annotations def _ad_convert_type(self, value, options=None): @@ -411,7 +380,104 @@ def __deepcopy__(self, memodict={}): return self.copy(deepcopy=True) +<<<<<<< Updated upstream class CofunctionMixin(FunctionMixin): def _ad_dot(self, other): return assemble(firedrake.action(self, other)) +======= +class RieszMap: + """Return a map from a dual to a primal function space. + + A `RieszMap` can be called on a `Cofunction` in the appropriate space + to yield the `Function` which is the Riesz representer under the. + + Parameters + ---------- + function_space_or_inner_product: FunctionSpace or DualSpace or Function or Cofunction or ufl.Form + The space from which to map, or a bilinear form defining an + inner product. + sobolev_space: String or ufl.sobolevspace.SobolevSpace. + Used to determine the inner product. + bcs: DirichletBC or list of DirichletBC + Boundary conditions to apply to the Riesz map. + solver_options: dict + A dictionary of PETSc options to be passed to the solver. + """ + + def __init__(self, function_space_or_inner_product=None, + sobolev_space=ufl.L2, *, bcs=None, solver_options=None): + if isinstance(function_space_or_inner_product, ufl.Form): + args = ufl.algorithms.extract_arguments( + function_space_or_inner_product + ) + if len(args) != 2: + raise ValueError(f"inner_product has arity {len(args)}, " + "should be 2.") + function_space = args[0].function_space() + inner_product = function_space_or_inner_product + else: + function_space = function_space_or_inner_product + if hasattr(function_space, "function_space"): + function_space = function_space.function_space() + if ufl.duals.is_dual(function_space): + function_space = function_space.dual() + + if str(sobolev_space) == "l2": + inner_product = "l2" + else: + from firedrake import TrialFunction, TestFunction + u = TrialFunction(function_space) + v = TestFunction(function_space) + inner_product = RieszMap._inner_product_form( + sobolev_space, u, v + ) + + self._function_space = function_space + self._inner_product = inner_product + self._bcs = bcs + self._solver_options = solver_options + + @staticmethod + def _inner_product_form(sobolev_space, u, v): + from firedrake import inner, dx, grad + inner_products = { + "L2": lambda u, v: inner(u, v)*dx, + "H1": lambda u, v: inner(u, v)*dx + inner(grad(u), grad(v))*dx + } + try: + return inner_products[str(sobolev_space)](u, v) + except KeyError: + raise ValueError("No inner product defined for Sobolev space " + f"{sobolev_space}.") + + @cached_property + def _solver(self): + from firedrake import (LinearVariationalSolver, + LinearVariationalProblem, Function, Cofunction) + rhs = Cofunction(self._function_space.dual()) + soln = Function(self._function_space) + lvp = LinearVariationalProblem(self._inner_product, rhs, soln, + bcs=self._bcs) + solver = LinearVariationalSolver( + lvp, solver_parameters=self._solver_options + ) + return solver.solve, rhs, soln + + def __call__(self, cofunction): + """Return the Riesz representer of a Cofunction.""" + from firedrake import Function + if cofunction.function_space().dual() != self._function_space: + raise ValueError("Function space mismatch in RieszMap.") + output = Function(self._function_space) + + if self._inner_product == "l2": + output.dat.data[:] = cofunction.dat.data[:] + else: + solve, rhs, soln = self._solver + rhs.assign(cofunction) + solve() + output = Function(self._function_space) + output.assign(soln) + return output +>>>>>>> Stashed changes diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index bca08c8213..d85b869add 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -16,8 +16,8 @@ class Cofunction(ufl.Cofunction, CofunctionMixin): r"""A :class:`Cofunction` represents a function on a dual space. - Like Functions, cofunctions are - represented as sums of basis functions: + + Like Functions, cofunctions are represented as sums of basis functions: .. math:: diff --git a/firedrake/ufl_expr.py b/firedrake/ufl_expr.py index 1f9a66df75..f7108a9cc7 100644 --- a/firedrake/ufl_expr.py +++ b/firedrake/ufl_expr.py @@ -45,6 +45,12 @@ def __init__(self, function_space, number, part=None): number, part=part) self._function_space = function_space + def arguments(self): + return (self,) + + def coefficients(self): + return () + @utils.cached_property def cell_node_map(self): return self.function_space().cell_node_map From c373c3120f76d0c840499c405ce51e57799ce90e Mon Sep 17 00:00:00 2001 From: David Ham Date: Mon, 1 Jul 2024 17:51:31 +0100 Subject: [PATCH 05/23] missing import --- firedrake/adjoint_utils/function.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 8ef0fa6696..f1a6d86738 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -380,12 +380,12 @@ def __deepcopy__(self, memodict={}): return self.copy(deepcopy=True) -<<<<<<< Updated upstream class CofunctionMixin(FunctionMixin): def _ad_dot(self, other): - return assemble(firedrake.action(self, other)) -======= + return firedrake.assemble(firedrake.action(self, other)) + + class RieszMap: """Return a map from a dual to a primal function space. @@ -480,4 +480,3 @@ def __call__(self, cofunction): output = Function(self._function_space) output.assign(soln) return output ->>>>>>> Stashed changes From 90cd52bf0a5f7122d6d26c020f06417edb4dd336 Mon Sep 17 00:00:00 2001 From: David Ham Date: Thu, 4 Jul 2024 11:27:20 +0100 Subject: [PATCH 06/23] fix l2 for mixed spaces --- firedrake/adjoint_utils/function.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index f1a6d86738..c17fbc28de 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -472,7 +472,8 @@ def __call__(self, cofunction): output = Function(self._function_space) if self._inner_product == "l2": - output.dat.data[:] = cofunction.dat.data[:] + for o, c in zip(output.subfunctions, cofunction.subfunctions): + o.dat.data[:] = c.dat.data[:] else: solve, rhs, soln = self._solver rhs.assign(cofunction) From fc45c48bb8c6397e90e56e57eaf996ad548fda44 Mon Sep 17 00:00:00 2001 From: "David A. Ham" Date: Thu, 4 Jul 2024 13:35:03 +0100 Subject: [PATCH 07/23] fix reisz map for primal case --- firedrake/adjoint_utils/function.py | 57 +++++++++++++++++++---------- firedrake/function.py | 9 ++--- 2 files changed, 41 insertions(+), 25 deletions(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index c17fbc28de..7bdb5cbc94 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -229,7 +229,7 @@ def _ad_convert_riesz(self, value, options=None): if value == 0.: # In adjoint-based differentiation, value == 0. arises only when # the functional is independent on the control variable. - return Function(V) + return Function(self.function_space()) if not isinstance(value, (Cofunction, Function)): raise TypeError("Expected a Cofunction or a Function") @@ -263,9 +263,7 @@ def _ad_restore_at_checkpoint(self, checkpoint): return checkpoint def _ad_will_add_as_dependency(self): - """Method called when the object is added as a Block dependency. - - """ + """Method called when the object is added as a Block dependency.""" with checkpoint_init_data(): super()._ad_will_add_as_dependency() @@ -273,7 +271,8 @@ def _ad_mul(self, other): from firedrake import Function r = Function(self.function_space()) - # `self` can be a Cofunction in which case only left multiplication with a scalar is allowed. + # `self` can be a Cofunction in which case only left multiplication + # with a scalar is allowed. r.assign(other * self) return r @@ -406,7 +405,7 @@ class RieszMap: """ def __init__(self, function_space_or_inner_product=None, - sobolev_space=ufl.L2, *, bcs=None, solver_options=None): + sobolev_space=ufl.L2, *, bcs=None, solver_options=None): if isinstance(function_space_or_inner_product, ufl.Form): args = ufl.algorithms.extract_arguments( function_space_or_inner_product @@ -464,20 +463,38 @@ def _solver(self): ) return solver.solve, rhs, soln - def __call__(self, cofunction): - """Return the Riesz representer of a Cofunction.""" - from firedrake import Function - if cofunction.function_space().dual() != self._function_space: - raise ValueError("Function space mismatch in RieszMap.") - output = Function(self._function_space) + def __call__(self, value): + """Return the Riesz representer of a Function or Cofunction.""" + from firedrake import Function, Cofunction - if self._inner_product == "l2": - for o, c in zip(output.subfunctions, cofunction.subfunctions): - o.dat.data[:] = c.dat.data[:] - else: - solve, rhs, soln = self._solver - rhs.assign(cofunction) - solve() + if ufl.duals.is_dual(value): + if value.function_space().dual() != self._function_space: + raise ValueError("Function space mismatch in RieszMap.") output = Function(self._function_space) - output.assign(soln) + + if self._inner_product == "l2": + for o, c in zip(output.subfunctions, value.subfunctions): + o.dat.data[:] = c.dat.data[:] + else: + solve, rhs, soln = self._solver + rhs.assign(value) + solve() + output = Function(self._function_space) + output.assign(soln) + elif ufl.duals.is_primal(value): + if value.function_space().dual() != self._function_space: + raise ValueError("Function space mismatch in RieszMap.") + + if self._inner_product == "l2": + output = Cofunction(self._function_space.dual()) + for o, c in zip(output.subfunctions, value.subfunctions): + o.dat.data[:] = c.dat.data[:] + else: + output = firedrake.assemble( + firedrake.action(self._inner_product, value) + ) + else: + raise ValueError( + f"Unable to ascertain if {value} is primal or dual." + ) return output diff --git a/firedrake/function.py b/firedrake/function.py index da4d264971..18c83539f9 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -21,7 +21,7 @@ from firedrake.cofunction import Cofunction from firedrake import utils from firedrake import vector -from firedrake.adjoint_utils import FunctionMixin +from firedrake.adjoint_utils import FunctionMixin, RieszMap from firedrake.petsc import PETSc @@ -502,11 +502,10 @@ def riesz_representation(self, riesz_map='L2'): if riesz_map == "l2": return Cofunction(V.dual(), val=self.dat) - elif riesz_map in ("L2", "H1"): - a = self._define_riesz_map_form(riesz_map, V) - return assemble(action(a, self)) + if riesz_map in ("L2", "H1"): + riesz_map = RieszMap(V, riesz_map) - elif callable(riesz_map): + if callable(riesz_map): return riesz_map(self) else: From 1c16bb5e579d3355b57e4ac23e7f3d84f9f51e4a Mon Sep 17 00:00:00 2001 From: "David A. Ham" Date: Thu, 4 Jul 2024 15:11:07 +0100 Subject: [PATCH 08/23] riesz_representation --- firedrake/cofunction.py | 45 ++++++++++++++++++++--------------------- firedrake/function.py | 43 ++++++++++++++++++--------------------- 2 files changed, 42 insertions(+), 46 deletions(-) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index d85b869add..f58888274c 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -8,7 +8,7 @@ import firedrake.functionspaceimpl as functionspaceimpl from firedrake import utils, vector, ufl_expr from firedrake.utils import ScalarType -from firedrake.adjoint_utils.function import CofunctionMixin +from firedrake.adjoint_utils.function import CofunctionMixin, RieszMap from firedrake.adjoint_utils.checkpointing import DelegatedFunctionCheckpoint from firedrake.adjoint_utils.blocks.function import CofunctionAssignBlock from firedrake.petsc import PETSc @@ -228,37 +228,36 @@ def assign(self, expr, subset=None, expr_from_assemble=False): raise ValueError('Cannot assign %s' % expr) - def riesz_representation(self, riesz_map='L2', **solver_options): - """Return the Riesz representation of this :class:`Cofunction` with respect to the given Riesz map. + def riesz_representation(self, riesz_map='L2', *, bcs=None, + solver_options=None): + """Return the Riesz representation of this :class:`Cofunction`. - Example: For a L2 Riesz map, the Riesz representation is obtained by solving - the linear system ``Mx = self``, where M is the L2 mass matrix, i.e. M = - with u and v trial and test functions, respectively. + Example: For a L2 Riesz map, the Riesz representation is obtained by + solving the linear system ``Mx = self``, where M is the L2 mass matrix, + i.e. M = with u and v trial and test functions, respectively. Parameters ---------- - riesz_map : str or collections.abc.Callable - The Riesz map to use (`l2`, `L2`, or `H1`). This can also be a callable. - solver_options : dict - Solver options to pass to the linear solver: - - solver_parameters: optional solver parameters. - - nullspace: an optional :class:`.VectorSpaceBasis` (or :class:`.MixedVectorSpaceBasis`) - spanning the null space of the operator. - - transpose_nullspace: as for the nullspace, but used to make the right hand side consistent. - - near_nullspace: as for the nullspace, but used to add the near nullspace. - - options_prefix: an optional prefix used to distinguish PETSc options. - If not provided a unique prefix will be created. - Use this option if you want to pass options to the solver from the command line - in addition to through the ``solver_parameters`` dict. + riesz_map : str or ufl.sobolevspace.SobolevSpace or + collections.abc.Callable + The Riesz map to use (`l2`, `L2`, or `H1`). This can also be a + callable. + bcs: DirichletBC or list of DirichletBC + Boundary conditions to apply to the Riesz map. + solver_options: dict + A dictionary of PETSc options to be passed to the solver. Returns ------- firedrake.function.Function - Riesz representation of this :class:`Cofunction` with respect to the given Riesz map. + Riesz representation of this :class:`Cofunction` with respect to + the given Riesz map. """ - return self._ad_convert_riesz(self, options={"function_space": self.function_space().dual(), - "riesz_representation": riesz_map, - "solver_options": solver_options}) + if not callable(riesz_map): + riesz_map = RieszMap(self.function_space(), riesz_map, bcs=bcs, + solver_options=solver_options) + + return riesz_map(self) @CofunctionMixin._ad_annotate_iadd @utils.known_pyop2_safe diff --git a/firedrake/function.py b/firedrake/function.py index 18c83539f9..94ebdc7cb3 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -478,39 +478,36 @@ def assign(self, expr, subset=None): Assigner(self, expr, subset).assign() return self - def riesz_representation(self, riesz_map='L2'): - """Return the Riesz representation of this :class:`Function` with respect to the given Riesz map. + def riesz_representation(self, riesz_map='L2', bcs=None, + solver_options=None): + """Return the Riesz representation of this :class:`Function`. - Example: For a L2 Riesz map, the Riesz representation is obtained by taking the action - of ``M`` on ``self``, where M is the L2 mass matrix, i.e. M = - with u and v trial and test functions, respectively. + Example: For a L2 Riesz map, the Riesz representation is obtained by + taking the action of ``M`` on ``self``, where M is the L2 mass matrix, + i.e. M = with u and v trial and test functions, respectively. Parameters ---------- - riesz_map : str or collections.abc.Callable - The Riesz map to use (`l2`, `L2`, or `H1`). This can also be a callable. + riesz_map : str or ufl.sobolevspace.SobolevSpace or + collections.abc.Callable + The Riesz map to use (`l2`, `L2`, or `H1`). This can also be a + callable which applies the Riesz map. + bcs: DirichletBC or list of DirichletBC + Boundary conditions to apply to the Riesz map. + solver_options: dict + A dictionary of PETSc options to be passed to the solver. Returns ------- firedrake.cofunction.Cofunction - Riesz representation of this :class:`Function` with respect to the given Riesz map. + Riesz representation of this :class:`Function` with respect to the + given Riesz map. """ - from firedrake.ufl_expr import action - from firedrake.assemble import assemble + if not callable(riesz_map): + riesz_map = RieszMap(self.function_space(), riesz_map, bcs=bcs, + solver_options=solver_options) - V = self.function_space() - if riesz_map == "l2": - return Cofunction(V.dual(), val=self.dat) - - if riesz_map in ("L2", "H1"): - riesz_map = RieszMap(V, riesz_map) - - if callable(riesz_map): - return riesz_map(self) - - else: - raise NotImplementedError( - "Unknown Riesz representation %s" % riesz_map) + return riesz_map(self) @FunctionMixin._ad_annotate_iadd def __iadd__(self, expr): From 42f974cbe941161489cc6f3217588cafce6b77c0 Mon Sep 17 00:00:00 2001 From: "David A. Ham" Date: Thu, 4 Jul 2024 15:18:29 +0100 Subject: [PATCH 09/23] Update firedrake/adjoint_utils/blocks/solving.py --- firedrake/adjoint_utils/blocks/solving.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index 2028bc2684..6412f5b2e5 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -794,7 +794,7 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, dFdm = replace(dFdm, replace_map) if isinstance(dFdm, firedrake.Argument): - # Corner case. Should be fixed more permanenty upstream in UFL. + # Corner case. Should be fixed more permanently upstream in UFL. dFdm = ufl.Action(dFdm, adj_sol) else: dFdm = dFdm * adj_sol From 8b93fe4aba6e37035262bf2c49c126548e1d99a6 Mon Sep 17 00:00:00 2001 From: "David A. Ham" Date: Thu, 4 Jul 2024 16:02:58 +0100 Subject: [PATCH 10/23] rearrange to pass documentation --- firedrake/adjoint_utils/function.py | 129 +--------------------------- firedrake/cofunction.py | 118 ++++++++++++++++++++++++- firedrake/function.py | 4 +- 3 files changed, 123 insertions(+), 128 deletions(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 7bdb5cbc94..1e3301e63d 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -1,4 +1,4 @@ -from functools import wraps, cached_property +from functools import wraps import ufl from ufl.domain import extract_unique_domain from pyadjoint.overloaded_type import create_overloaded_object, FloatingType @@ -221,7 +221,7 @@ def _ad_create_checkpoint(self): return self.copy(deepcopy=True) def _ad_convert_riesz(self, value, options=None): - from firedrake import Function, Cofunction + from firedrake import Function options = {} if options is None else options riesz_representation = options.get("riesz_representation", "L2") @@ -231,14 +231,8 @@ def _ad_convert_riesz(self, value, options=None): # the functional is independent on the control variable. return Function(self.function_space()) - if not isinstance(value, (Cofunction, Function)): - raise TypeError("Expected a Cofunction or a Function") - - if callable(riesz_representation): - return riesz_representation(value) - - return RieszMap(self, riesz_representation, - solver_options=solver_options)(value) + return value.riesz_representation(riesz_map=riesz_representation, + solver_options=solver_options) @no_annotations def _ad_convert_type(self, value, options=None): @@ -383,118 +377,3 @@ class CofunctionMixin(FunctionMixin): def _ad_dot(self, other): return firedrake.assemble(firedrake.action(self, other)) - - -class RieszMap: - """Return a map from a dual to a primal function space. - - A `RieszMap` can be called on a `Cofunction` in the appropriate space - to yield the `Function` which is the Riesz representer under the. - - Parameters - ---------- - function_space_or_inner_product: FunctionSpace or DualSpace or Function or Cofunction or ufl.Form - The space from which to map, or a bilinear form defining an - inner product. - sobolev_space: String or ufl.sobolevspace.SobolevSpace. - Used to determine the inner product. - bcs: DirichletBC or list of DirichletBC - Boundary conditions to apply to the Riesz map. - solver_options: dict - A dictionary of PETSc options to be passed to the solver. - """ - - def __init__(self, function_space_or_inner_product=None, - sobolev_space=ufl.L2, *, bcs=None, solver_options=None): - if isinstance(function_space_or_inner_product, ufl.Form): - args = ufl.algorithms.extract_arguments( - function_space_or_inner_product - ) - if len(args) != 2: - raise ValueError(f"inner_product has arity {len(args)}, " - "should be 2.") - function_space = args[0].function_space() - inner_product = function_space_or_inner_product - else: - function_space = function_space_or_inner_product - if hasattr(function_space, "function_space"): - function_space = function_space.function_space() - if ufl.duals.is_dual(function_space): - function_space = function_space.dual() - - if str(sobolev_space) == "l2": - inner_product = "l2" - else: - from firedrake import TrialFunction, TestFunction - u = TrialFunction(function_space) - v = TestFunction(function_space) - inner_product = RieszMap._inner_product_form( - sobolev_space, u, v - ) - - self._function_space = function_space - self._inner_product = inner_product - self._bcs = bcs - self._solver_options = solver_options - - @staticmethod - def _inner_product_form(sobolev_space, u, v): - from firedrake import inner, dx, grad - inner_products = { - "L2": lambda u, v: inner(u, v)*dx, - "H1": lambda u, v: inner(u, v)*dx + inner(grad(u), grad(v))*dx - } - try: - return inner_products[str(sobolev_space)](u, v) - except KeyError: - raise ValueError("No inner product defined for Sobolev space " - f"{sobolev_space}.") - - @cached_property - def _solver(self): - from firedrake import (LinearVariationalSolver, - LinearVariationalProblem, Function, Cofunction) - rhs = Cofunction(self._function_space.dual()) - soln = Function(self._function_space) - lvp = LinearVariationalProblem(self._inner_product, rhs, soln, - bcs=self._bcs) - solver = LinearVariationalSolver( - lvp, solver_parameters=self._solver_options - ) - return solver.solve, rhs, soln - - def __call__(self, value): - """Return the Riesz representer of a Function or Cofunction.""" - from firedrake import Function, Cofunction - - if ufl.duals.is_dual(value): - if value.function_space().dual() != self._function_space: - raise ValueError("Function space mismatch in RieszMap.") - output = Function(self._function_space) - - if self._inner_product == "l2": - for o, c in zip(output.subfunctions, value.subfunctions): - o.dat.data[:] = c.dat.data[:] - else: - solve, rhs, soln = self._solver - rhs.assign(value) - solve() - output = Function(self._function_space) - output.assign(soln) - elif ufl.duals.is_primal(value): - if value.function_space().dual() != self._function_space: - raise ValueError("Function space mismatch in RieszMap.") - - if self._inner_product == "l2": - output = Cofunction(self._function_space.dual()) - for o, c in zip(output.subfunctions, value.subfunctions): - o.dat.data[:] = c.dat.data[:] - else: - output = firedrake.assemble( - firedrake.action(self._inner_product, value) - ) - else: - raise ValueError( - f"Unable to ascertain if {value} is primal or dual." - ) - return output diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index f58888274c..cd4b150218 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -1,3 +1,4 @@ +from functools import cached_property import numpy as np import ufl @@ -8,7 +9,7 @@ import firedrake.functionspaceimpl as functionspaceimpl from firedrake import utils, vector, ufl_expr from firedrake.utils import ScalarType -from firedrake.adjoint_utils.function import CofunctionMixin, RieszMap +from firedrake.adjoint_utils.function import CofunctionMixin from firedrake.adjoint_utils.checkpointing import DelegatedFunctionCheckpoint from firedrake.adjoint_utils.blocks.function import CofunctionAssignBlock from firedrake.petsc import PETSc @@ -359,3 +360,118 @@ def __str__(self): def cell_node_map(self): return self.function_space().cell_node_map() + + +class RieszMap: + """Return a map from a dual to a primal function space. + + A `RieszMap` can be called on a `Cofunction` in the appropriate space + to yield the `Function` which is the Riesz representer under the. + + Parameters + ---------- + function_space_or_inner_product: FunctionSpace or Function or Cofunction or ufl.Form + The space from which to map, or a bilinear form defining an + inner product. + sobolev_space: str or ufl.sobolevspace.SobolevSpace. + Used to determine the inner product. + bcs: DirichletBC or list of DirichletBC + Boundary conditions to apply to the Riesz map. + solver_options: dict + A dictionary of PETSc options to be passed to the solver. + """ + + def __init__(self, function_space_or_inner_product=None, + sobolev_space=ufl.L2, *, bcs=None, solver_options=None): + if isinstance(function_space_or_inner_product, ufl.Form): + args = ufl.algorithms.extract_arguments( + function_space_or_inner_product + ) + if len(args) != 2: + raise ValueError(f"inner_product has arity {len(args)}, " + "should be 2.") + function_space = args[0].function_space() + inner_product = function_space_or_inner_product + else: + function_space = function_space_or_inner_product + if hasattr(function_space, "function_space"): + function_space = function_space.function_space() + if ufl.duals.is_dual(function_space): + function_space = function_space.dual() + + if str(sobolev_space) == "l2": + inner_product = "l2" + else: + from firedrake import TrialFunction, TestFunction + u = TrialFunction(function_space) + v = TestFunction(function_space) + inner_product = RieszMap._inner_product_form( + sobolev_space, u, v + ) + + self._function_space = function_space + self._inner_product = inner_product + self._bcs = bcs + self._solver_options = solver_options + + @staticmethod + def _inner_product_form(sobolev_space, u, v): + from firedrake import inner, dx, grad + inner_products = { + "L2": lambda u, v: inner(u, v)*dx, + "H1": lambda u, v: inner(u, v)*dx + inner(grad(u), grad(v))*dx + } + try: + return inner_products[str(sobolev_space)](u, v) + except KeyError: + raise ValueError("No inner product defined for Sobolev space " + f"{sobolev_space}.") + + @cached_property + def _solver(self): + from firedrake import (LinearVariationalSolver, + LinearVariationalProblem, Function, Cofunction) + rhs = Cofunction(self._function_space.dual()) + soln = Function(self._function_space) + lvp = LinearVariationalProblem(self._inner_product, rhs, soln, + bcs=self._bcs) + solver = LinearVariationalSolver( + lvp, solver_parameters=self._solver_options + ) + return solver.solve, rhs, soln + + def __call__(self, value): + """Return the Riesz representer of a Function or Cofunction.""" + from firedrake import Function, Cofunction + + if ufl.duals.is_dual(value): + if value.function_space().dual() != self._function_space: + raise ValueError("Function space mismatch in RieszMap.") + output = Function(self._function_space) + + if self._inner_product == "l2": + for o, c in zip(output.subfunctions, value.subfunctions): + o.dat.data[:] = c.dat.data[:] + else: + solve, rhs, soln = self._solver + rhs.assign(value) + solve() + output = Function(self._function_space) + output.assign(soln) + elif ufl.duals.is_primal(value): + if value.function_space().dual() != self._function_space: + raise ValueError("Function space mismatch in RieszMap.") + + if self._inner_product == "l2": + output = Cofunction(self._function_space.dual()) + for o, c in zip(output.subfunctions, value.subfunctions): + o.dat.data[:] = c.dat.data[:] + else: + output = firedrake.assemble( + firedrake.action(self._inner_product, value) + ) + else: + raise ValueError( + f"Unable to ascertain if {value} is primal or dual." + ) + return output diff --git a/firedrake/function.py b/firedrake/function.py index 94ebdc7cb3..fad5502e7f 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -18,10 +18,10 @@ from firedrake.utils import ScalarType, IntType, as_ctypes from firedrake import functionspaceimpl -from firedrake.cofunction import Cofunction +from firedrake.cofunction import Cofunction, RieszMap from firedrake import utils from firedrake import vector -from firedrake.adjoint_utils import FunctionMixin, RieszMap +from firedrake.adjoint_utils import FunctionMixin from firedrake.petsc import PETSc From 85a8267047ad3f6113658222be56f26ecb6f91e9 Mon Sep 17 00:00:00 2001 From: "David A. Ham" Date: Thu, 4 Jul 2024 16:18:41 +0100 Subject: [PATCH 11/23] fix solver options --- firedrake/cofunction.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index cd4b150218..c7dbb80549 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -412,7 +412,7 @@ def __init__(self, function_space_or_inner_product=None, self._function_space = function_space self._inner_product = inner_product self._bcs = bcs - self._solver_options = solver_options + self._solver_options = solver_options or {} @staticmethod def _inner_product_form(sobolev_space, u, v): From 6768c9ffe0e9a309ea808808690f714c786564ad Mon Sep 17 00:00:00 2001 From: David Ham Date: Sat, 6 Jul 2024 20:54:39 +0100 Subject: [PATCH 12/23] make solve string more robust --- firedrake/adjoint_utils/blocks/solving.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/firedrake/adjoint_utils/blocks/solving.py b/firedrake/adjoint_utils/blocks/solving.py index 6412f5b2e5..4c4020a249 100644 --- a/firedrake/adjoint_utils/blocks/solving.py +++ b/firedrake/adjoint_utils/blocks/solving.py @@ -86,8 +86,15 @@ def _init_solver_parameters(self, args, kwargs): self.assemble_kwargs = {} def __str__(self): - return "solve({} = {})".format(ufl2unicode(self.lhs), - ufl2unicode(self.rhs)) + try: + lhs_string = ufl2unicode(self.lhs) + except AttributeError: + lhs_string = str(self.lhs) + try: + rhs_string = ufl2unicode(self.rhs) + except AttributeError: + rhs_string = str(self.rhs) + return "solve({} = {})".format(lhs_string, rhs_string) def _create_F_form(self): # Process the equation forms, replacing values with checkpoints, From f98e837743190282742480624a153d83d0a41872 Mon Sep 17 00:00:00 2001 From: David Ham Date: Mon, 8 Jul 2024 17:36:01 +0100 Subject: [PATCH 13/23] docstrings --- firedrake/cofunction.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index c7dbb80549..689269c0aa 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -363,16 +363,18 @@ def cell_node_map(self): class RieszMap: - """Return a map from a dual to a primal function space. + """Return a map between dual and primal function spaces. - A `RieszMap` can be called on a `Cofunction` in the appropriate space - to yield the `Function` which is the Riesz representer under the. + A `RieszMap` can be called on a `Cofunction` in the appropriate space to + yield the `Function` which is the Riesz representer under the given inner + product. Conversely, it can be called on a `Function` to apply the given + inner product and return a `Cofunction`. Parameters ---------- - function_space_or_inner_product: FunctionSpace or Function or Cofunction or ufl.Form - The space from which to map, or a bilinear form defining an - inner product. + function_space_or_inner_product: FunctionSpace or ufl.Form + The space from which to map, or a bilinear form defining an inner + product. sobolev_space: str or ufl.sobolevspace.SobolevSpace. Used to determine the inner product. bcs: DirichletBC or list of DirichletBC From 1d8e58a6cf1828c91c14a14a2d224ebcba99c366 Mon Sep 17 00:00:00 2001 From: David Ham Date: Tue, 17 Dec 2024 15:14:59 +0000 Subject: [PATCH 14/23] remove merged branch reference --- .github/workflows/build.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index acebd51af8..909db6990a 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -83,7 +83,6 @@ jobs: --install defcon \ --install gadopt \ --install asQ \ - --package-branch ufl dham/cofunction_is_terminal \ || (cat firedrake-install.log && /bin/false) - name: Install test dependencies run: | From cc402aa47b978d47d566af55b8c81e56104cb6ff Mon Sep 17 00:00:00 2001 From: David Ham Date: Tue, 17 Dec 2024 15:20:25 +0000 Subject: [PATCH 15/23] lint --- tests/firedrake/regression/test_adjoint_operators.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/firedrake/regression/test_adjoint_operators.py b/tests/firedrake/regression/test_adjoint_operators.py index d0abec9a0b..8d4c0ab796 100644 --- a/tests/firedrake/regression/test_adjoint_operators.py +++ b/tests/firedrake/regression/test_adjoint_operators.py @@ -1002,6 +1002,7 @@ def test_cofunction_assign_functional(): f.assign(2.0) assert np.allclose(Jhat(f), 2.0) + @pytest.mark.skipcomplex # Taping for complex-valued 0-forms not yet done def test_bdy_control(): # Test for the case the boundary condition is a control for a From cde10edd7112d05d63f1ca17cf3487d534906359 Mon Sep 17 00:00:00 2001 From: David Ham Date: Tue, 17 Dec 2024 16:20:00 +0000 Subject: [PATCH 16/23] Remove unused solver options --- firedrake/cofunction.py | 33 +++++++++++++++++++++++---------- firedrake/function.py | 8 ++------ 2 files changed, 25 insertions(+), 16 deletions(-) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index 689269c0aa..f3981b7b4c 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -230,7 +230,8 @@ def assign(self, expr, subset=None, expr_from_assemble=False): raise ValueError('Cannot assign %s' % expr) def riesz_representation(self, riesz_map='L2', *, bcs=None, - solver_options=None): + solver_parameters=None, + form_compiler_parameters=None): """Return the Riesz representation of this :class:`Cofunction`. Example: For a L2 Riesz map, the Riesz representation is obtained by @@ -245,8 +246,11 @@ def riesz_representation(self, riesz_map='L2', *, bcs=None, callable. bcs: DirichletBC or list of DirichletBC Boundary conditions to apply to the Riesz map. - solver_options: dict + solver_parameters: dict A dictionary of PETSc options to be passed to the solver. + form_compiler_parameters: dict + A dictionary of form compiler parameters to be passed to the + variational problem that solves for the Riesz map. Returns ------- @@ -255,8 +259,11 @@ def riesz_representation(self, riesz_map='L2', *, bcs=None, the given Riesz map. """ if not callable(riesz_map): - riesz_map = RieszMap(self.function_space(), riesz_map, bcs=bcs, - solver_options=solver_options) + riesz_map = RieszMap( + self.function_space(), riesz_map, bcs=bcs, + solver_parameters=solver_parameters, + form_compiler_parameters=form_compiler_parameters + ) return riesz_map(self) @@ -379,12 +386,16 @@ class RieszMap: Used to determine the inner product. bcs: DirichletBC or list of DirichletBC Boundary conditions to apply to the Riesz map. - solver_options: dict + solver_parameters: dict A dictionary of PETSc options to be passed to the solver. + form_compiler_parameters: dict + A dictionary of form compiler parameters to be passed to the + variational problem that solves for the Riesz map. """ def __init__(self, function_space_or_inner_product=None, - sobolev_space=ufl.L2, *, bcs=None, solver_options=None): + sobolev_space=ufl.L2, *, bcs=None, solver_parameters=None, + form_compiler_parameters=None): if isinstance(function_space_or_inner_product, ufl.Form): args = ufl.algorithms.extract_arguments( function_space_or_inner_product @@ -414,7 +425,8 @@ def __init__(self, function_space_or_inner_product=None, self._function_space = function_space self._inner_product = inner_product self._bcs = bcs - self._solver_options = solver_options or {} + self._solver_parameters = solver_parameters or {} + self._form_compiler_parameters = form_compiler_parameters or {} @staticmethod def _inner_product_form(sobolev_space, u, v): @@ -435,10 +447,11 @@ def _solver(self): LinearVariationalProblem, Function, Cofunction) rhs = Cofunction(self._function_space.dual()) soln = Function(self._function_space) - lvp = LinearVariationalProblem(self._inner_product, rhs, soln, - bcs=self._bcs) + lvp = LinearVariationalProblem( + self._inner_product, rhs, soln, bcs=self._bcs, restrict=True, + form_compiler_parameters=self._form_compiler_parameters) solver = LinearVariationalSolver( - lvp, solver_parameters=self._solver_options + lvp, solver_parameters=self._solver_parameters ) return solver.solve, rhs, soln diff --git a/firedrake/function.py b/firedrake/function.py index fad5502e7f..d296a464f6 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -478,8 +478,7 @@ def assign(self, expr, subset=None): Assigner(self, expr, subset).assign() return self - def riesz_representation(self, riesz_map='L2', bcs=None, - solver_options=None): + def riesz_representation(self, riesz_map='L2', bcs=None): """Return the Riesz representation of this :class:`Function`. Example: For a L2 Riesz map, the Riesz representation is obtained by @@ -494,8 +493,6 @@ def riesz_representation(self, riesz_map='L2', bcs=None, callable which applies the Riesz map. bcs: DirichletBC or list of DirichletBC Boundary conditions to apply to the Riesz map. - solver_options: dict - A dictionary of PETSc options to be passed to the solver. Returns ------- @@ -504,8 +501,7 @@ def riesz_representation(self, riesz_map='L2', bcs=None, given Riesz map. """ if not callable(riesz_map): - riesz_map = RieszMap(self.function_space(), riesz_map, bcs=bcs, - solver_options=solver_options) + riesz_map = RieszMap(self.function_space(), riesz_map, bcs=bcs) return riesz_map(self) From 8f86232b0c7548cbd6a0c6a8504661de3f3da0da Mon Sep 17 00:00:00 2001 From: David Ham Date: Wed, 18 Dec 2024 10:20:53 +0000 Subject: [PATCH 17/23] wind back API changes --- firedrake/cofunction.py | 6 +++--- firedrake/function.py | 6 ++---- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/firedrake/cofunction.py b/firedrake/cofunction.py index f3981b7b4c..6e8bebe8bd 100644 --- a/firedrake/cofunction.py +++ b/firedrake/cofunction.py @@ -230,7 +230,7 @@ def assign(self, expr, subset=None, expr_from_assemble=False): raise ValueError('Cannot assign %s' % expr) def riesz_representation(self, riesz_map='L2', *, bcs=None, - solver_parameters=None, + solver_options=None, form_compiler_parameters=None): """Return the Riesz representation of this :class:`Cofunction`. @@ -246,7 +246,7 @@ def riesz_representation(self, riesz_map='L2', *, bcs=None, callable. bcs: DirichletBC or list of DirichletBC Boundary conditions to apply to the Riesz map. - solver_parameters: dict + solver_options: dict A dictionary of PETSc options to be passed to the solver. form_compiler_parameters: dict A dictionary of form compiler parameters to be passed to the @@ -261,7 +261,7 @@ def riesz_representation(self, riesz_map='L2', *, bcs=None, if not callable(riesz_map): riesz_map = RieszMap( self.function_space(), riesz_map, bcs=bcs, - solver_parameters=solver_parameters, + solver_parameters=solver_options, form_compiler_parameters=form_compiler_parameters ) diff --git a/firedrake/function.py b/firedrake/function.py index d296a464f6..5aae0fe561 100644 --- a/firedrake/function.py +++ b/firedrake/function.py @@ -478,7 +478,7 @@ def assign(self, expr, subset=None): Assigner(self, expr, subset).assign() return self - def riesz_representation(self, riesz_map='L2', bcs=None): + def riesz_representation(self, riesz_map='L2'): """Return the Riesz representation of this :class:`Function`. Example: For a L2 Riesz map, the Riesz representation is obtained by @@ -491,8 +491,6 @@ def riesz_representation(self, riesz_map='L2', bcs=None): collections.abc.Callable The Riesz map to use (`l2`, `L2`, or `H1`). This can also be a callable which applies the Riesz map. - bcs: DirichletBC or list of DirichletBC - Boundary conditions to apply to the Riesz map. Returns ------- @@ -501,7 +499,7 @@ def riesz_representation(self, riesz_map='L2', bcs=None): given Riesz map. """ if not callable(riesz_map): - riesz_map = RieszMap(self.function_space(), riesz_map, bcs=bcs) + riesz_map = RieszMap(self.function_space(), riesz_map) return riesz_map(self) From 720ad1dfdab1bfe3e37a4abf96236945545f50d9 Mon Sep 17 00:00:00 2001 From: David Ham Date: Wed, 18 Dec 2024 12:11:02 +0000 Subject: [PATCH 18/23] test for correct annotation of cofunction assign --- .../regression/test_adjoint_operators.py | 36 +++++++++++-------- 1 file changed, 21 insertions(+), 15 deletions(-) diff --git a/tests/firedrake/regression/test_adjoint_operators.py b/tests/firedrake/regression/test_adjoint_operators.py index 8d4c0ab796..479fd18a92 100644 --- a/tests/firedrake/regression/test_adjoint_operators.py +++ b/tests/firedrake/regression/test_adjoint_operators.py @@ -932,21 +932,27 @@ def test_riesz_representation_for_adjoints(): dJdu_function_L2 = Function(space) solve(a == dJdu_cofunction, dJdu_function_L2) - dJdu_none = rf.derivative(options={"riesz_representation": None}) - dJdu_l2 = rf.derivative(options={"riesz_representation": "l2"}) - dJdu_H1 = rf.derivative(options={"riesz_representation": "H1"}) - dJdu_L2 = rf.derivative(options={"riesz_representation": "L2"}) - dJdu_default_L2 = rf.derivative() - assert ( - isinstance(dJdu_none, Cofunction) and isinstance(dJdu_function_l2, Function) - and isinstance(dJdu_H1, Function) and isinstance(dJdu_default_L2, Function) - and isinstance(dJdu_L2, Function) - and np.allclose(dJdu_none.dat.data, dJdu_cofunction.dat.data) - and np.allclose(dJdu_l2.dat.data, dJdu_function_l2.dat.data) - and np.allclose(dJdu_H1.dat.data, dJdu_function_H1.dat.data) - and np.allclose(dJdu_default_L2.dat.data, dJdu_function_L2.dat.data) - and np.allclose(dJdu_L2.dat.data, dJdu_function_L2.dat.data) + dJdu_none = ReducedFunctional(J, Control(f)).derivative() + dJdu_l2 = ReducedFunctional(J, Control(f, riesz_map="l2")).derivative( + apply_riesz=True ) + dJdu_H1 = ReducedFunctional(J, Control(f, riesz_map="H1")).derivative( + apply_riesz=True + ) + dJdu_L2 = ReducedFunctional(J, Control(f, riesz_map="L2")).derivative( + apply_riesz=True + ) + dJdu_default_L2 = ReducedFunctional(J, Control(f)).derivative( + apply_riesz=True + ) + assert isinstance(dJdu_none, Cofunction) and isinstance(dJdu_function_l2, Function) + assert isinstance(dJdu_H1, Function) and isinstance(dJdu_default_L2, Function) + assert isinstance(dJdu_L2, Function) + assert np.allclose(dJdu_none.dat.data, dJdu_cofunction.dat.data) + assert np.allclose(dJdu_l2.dat.data, dJdu_function_l2.dat.data) + assert np.allclose(dJdu_H1.dat.data, dJdu_function_H1.dat.data) + assert np.allclose(dJdu_default_L2.dat.data, dJdu_function_L2.dat.data) + assert np.allclose(dJdu_L2.dat.data, dJdu_function_L2.dat.data) @pytest.mark.skipcomplex @@ -998,7 +1004,7 @@ def test_cofunction_assign_functional(): cof2.assign(cof) # Test is checking that this is taped. J = assemble(action(cof2, f2)) Jhat = ReducedFunctional(J, Control(f)) - assert np.allclose(float(Jhat.derivative()), 1.0) + assert np.allclose(float(Jhat.derivative(apply_riesz=True)), 1.0) f.assign(2.0) assert np.allclose(Jhat(f), 2.0) From 674e4bb882cc5e905c01221827fa958ef9b6d239 Mon Sep 17 00:00:00 2001 From: David Ham Date: Wed, 18 Dec 2024 12:16:22 +0000 Subject: [PATCH 19/23] overloadedtype interface changes --- firedrake/adjoint_utils/function.py | 34 ++++--------------- .../regression/test_adjoint_operators.py | 8 ++--- 2 files changed, 11 insertions(+), 31 deletions(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 1e3301e63d..6b78163ed7 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -220,35 +220,15 @@ def _ad_create_checkpoint(self): else: return self.copy(deepcopy=True) - def _ad_convert_riesz(self, value, options=None): - from firedrake import Function + def _ad_convert_riesz(self, value, riesz_map=None): + return value.riesz_representation(riesz_map=riesz_map or "L2") - options = {} if options is None else options - riesz_representation = options.get("riesz_representation", "L2") - solver_options = options.get("solver_options", {}) - if value == 0.: - # In adjoint-based differentiation, value == 0. arises only when - # the functional is independent on the control variable. - return Function(self.function_space()) - - return value.riesz_representation(riesz_map=riesz_representation, - solver_options=solver_options) - - @no_annotations - def _ad_convert_type(self, value, options=None): - # `_ad_convert_type` is not annotated, unlike `_ad_convert_riesz` - options = {} if options is None else options.copy() - options.setdefault("riesz_representation", "L2") - if options["riesz_representation"] is None: - if value == 0.: - # In adjoint-based differentiation, value == 0. arises only when - # the functional is independent on the control variable. - V = options.get("function_space", self.function_space()) - return firedrake.Cofunction(V.dual()) - else: - return value + def _ad_init_zero(self, dual=False): + from firedrake import Function, Cofunction + if dual: + return Cofunction(self.function_space().dual()) else: - return self._ad_convert_riesz(value, options=options) + return Function(self.function_space()) def _ad_restore_at_checkpoint(self, checkpoint): if isinstance(checkpoint, CheckpointBase): diff --git a/tests/firedrake/regression/test_adjoint_operators.py b/tests/firedrake/regression/test_adjoint_operators.py index 479fd18a92..de1252be19 100644 --- a/tests/firedrake/regression/test_adjoint_operators.py +++ b/tests/firedrake/regression/test_adjoint_operators.py @@ -864,10 +864,10 @@ def test_assign_zero_cofunction(): J = assemble(((sol + Constant(1.0)) ** 2) * dx) # The zero assignment should break the tape and hence cause a zero # gradient. - grad_l2 = compute_gradient(J, Control(k), options={"riesz_representation": "l2"}) - grad_none = compute_gradient(J, Control(k), options={"riesz_representation": None}) - grad_h1 = compute_gradient(J, Control(k), options={"riesz_representation": "H1"}) - grad_L2 = compute_gradient(J, Control(k), options={"riesz_representation": "L2"}) + grad_l2 = compute_gradient(J, Control(k, riesz_map="l2"), apply_riesz=True) + grad_none = compute_gradient(J, Control(k)) + grad_h1 = compute_gradient(J, Control(k, riesz_map="H1"), apply_riesz=True) + grad_L2 = compute_gradient(J, Control(k, riesz_map="L2"), apply_riesz=True) assert isinstance(grad_l2, Function) and isinstance(grad_L2, Function) \ and isinstance(grad_h1, Function) assert isinstance(grad_none, Cofunction) From de14bf34ff7eef6e965459e60e8fdc2ef47e2119 Mon Sep 17 00:00:00 2001 From: David Ham Date: Wed, 18 Dec 2024 12:17:09 +0000 Subject: [PATCH 20/23] interface update in test --- tests/firedrake/regression/test_adjoint_operators.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/firedrake/regression/test_adjoint_operators.py b/tests/firedrake/regression/test_adjoint_operators.py index de1252be19..540932913b 100644 --- a/tests/firedrake/regression/test_adjoint_operators.py +++ b/tests/firedrake/regression/test_adjoint_operators.py @@ -913,7 +913,6 @@ def test_riesz_representation_for_adjoints(): space = FunctionSpace(mesh, "Lagrange", 1) f = Function(space).interpolate(SpatialCoordinate(mesh)[0]) J = assemble((f ** 2) * dx) - rf = ReducedFunctional(J, Control(f)) with stop_annotating(): v = TestFunction(space) u = TrialFunction(space) From cf57ae05b067de1774aecec9ab47f8a1209c4fbc Mon Sep 17 00:00:00 2001 From: David Ham Date: Wed, 18 Dec 2024 14:04:54 +0000 Subject: [PATCH 21/23] fix GlobalDataSet._apply_local_global_filter --- pyop2/types/dataset.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyop2/types/dataset.py b/pyop2/types/dataset.py index 087d9b091d..4a26cfc6d1 100644 --- a/pyop2/types/dataset.py +++ b/pyop2/types/dataset.py @@ -211,6 +211,7 @@ def __init__(self, global_): self._globalset = GlobalSet(comm=self.comm) self._name = "gdset_#x%x" % id(self) self._initialized = True + self._apply_local_global_filter = False @classmethod def _cache_key(cls, *args): From 55e11940c50dc98446ae995dfcfdd52021106753 Mon Sep 17 00:00:00 2001 From: David Ham Date: Wed, 18 Dec 2024 16:21:41 +0000 Subject: [PATCH 22/23] DROP BEFORE MERGE --- .github/workflows/build.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 909db6990a..aaf4239556 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -83,6 +83,7 @@ jobs: --install defcon \ --install gadopt \ --install asQ \ + --package-branch pyadjoint dham/abstract_reduced_functional || (cat firedrake-install.log && /bin/false) - name: Install test dependencies run: | From 112750e8d69548a348fbe46eff241738f4649453 Mon Sep 17 00:00:00 2001 From: David Ham Date: Wed, 18 Dec 2024 16:24:33 +0000 Subject: [PATCH 23/23] lint --- firedrake/adjoint_utils/function.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/firedrake/adjoint_utils/function.py b/firedrake/adjoint_utils/function.py index 6b78163ed7..f593bf605d 100644 --- a/firedrake/adjoint_utils/function.py +++ b/firedrake/adjoint_utils/function.py @@ -2,7 +2,7 @@ import ufl from ufl.domain import extract_unique_domain from pyadjoint.overloaded_type import create_overloaded_object, FloatingType -from pyadjoint.tape import annotate_tape, stop_annotating, get_working_tape, no_annotations +from pyadjoint.tape import annotate_tape, stop_annotating, get_working_tape from firedrake.adjoint_utils.blocks import FunctionAssignBlock, ProjectBlock, SubfunctionBlock, FunctionMergeBlock, SupermeshProjectBlock import firedrake from .checkpointing import disk_checkpointing, CheckpointFunction, \