From ace31ae43fc371b5bbdbfc5f04d0c5ee7b2c6162 Mon Sep 17 00:00:00 2001 From: nbouziani Date: Wed, 19 Oct 2022 13:21:16 +0100 Subject: [PATCH 1/7] Extend assemble block to encompass 1-forms --- dolfin_adjoint_common/blocks/assembly.py | 74 +++++++++++++++++++----- 1 file changed, 58 insertions(+), 16 deletions(-) diff --git a/dolfin_adjoint_common/blocks/assembly.py b/dolfin_adjoint_common/blocks/assembly.py index d9a84e40..a8aa7360 100644 --- a/dolfin_adjoint_common/blocks/assembly.py +++ b/dolfin_adjoint_common/blocks/assembly.py @@ -18,6 +18,41 @@ def __init__(self, form, ad_block_tag=None): def __str__(self): return f"assemble({ufl2unicode(self.form)})" + def compute_action_adjoint(self, adj_input, arity_form, form=None, c_rep=None, space=None, dform=None): + if arity_form == 0: + if dform is None: + dc = self.backend.TestFunction(space) + dform = self.backend.derivative(form, c_rep, dc) + dform_vector = self.compat.assemble_adjoint_value(dform) + return adj_input * dform_vector, dform + elif arity_form == 1: + if dform is None: + dc = self.backend.TrialFunction(space) + dform = self.backend.derivative(form, c_rep, dc) + # Get the Function + adj_input = adj_input.function + # Symbolic operators such as action/adjoint require derivatives to have been expanded beforehand. + # However, UFL doesn't support expanding coordinate derivatives of Coefficients in physical space, + # implying that we can't symbolically take the action/adjoint of the Jacobian for SpatialCoordinates. + # -> Workaround: Apply action/adjoint numerically (using PETSc). + if not isinstance(c_rep, self.backend.SpatialCoordinate): + # Symbolically compute: (dform/dc_rep)^* * adj_input + adj_output = self.backend.action(self.backend.adjoint(dform), adj_input) + adj_output = self.compat.assemble_adjoint_value(adj_output) + else: + # Get PETSc matrix + dform_mat = self.compat.assemble_adjoint_value(dform).petscmat + # Adjoint (hermitian transpose) + dform_mat.hermitianTranspose() + # Action + adj_output = self.backend.Function(space) + with adj_input.dat.vec_ro as v_vec: + with adj_output.dat.vec as res_vec: + dform_mat.mult(v_vec, res_vec) + return adj_output, dform + else: + raise ValueError('Forms with arity > 1 are not handled yet!') + def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies): replaced_coeffs = {} for block_variable in self.get_dependencies(): @@ -35,6 +70,9 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar c = block_variable.output c_rep = block_variable.saved_output + from ufl.algorithms.analysis import extract_arguments + arity_form = len(extract_arguments(form)) + if isinstance(c, self.compat.ExpressionType): # Create a FunctionSpace from self.form and Expression. # And then make a TestFunction from this space. @@ -47,17 +85,15 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar return [[adj_input * output, V]] if isinstance(c, self.backend.Function): - dc = self.backend.TestFunction(c.function_space()) + space = c.function_space() elif isinstance(c, self.backend.Constant): mesh = self.compat.extract_mesh_from_form(self.form) - dc = self.backend.TestFunction(c._ad_function_space(mesh)) + space = c._ad_function_space(mesh) elif isinstance(c, self.compat.MeshType): c_rep = self.backend.SpatialCoordinate(c_rep) - dc = self.backend.TestFunction(c._ad_function_space()) + space = c._ad_function_space() - dform = self.backend.derivative(form, c_rep, dc) - output = self.compat.assemble_adjoint_value(dform) - return adj_input * output + return self.compute_action_adjoint(adj_input, arity_form, form, c_rep, space)[0] def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs): return self.prepare_evaluate_adj(inputs, tlm_inputs, self.get_dependencies()) @@ -65,6 +101,9 @@ def prepare_evaluate_tlm(self, inputs, tlm_inputs, relevant_outputs): def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None): form = prepared dform = 0. + + from ufl.algorithms.analysis import extract_arguments + arity_form = len(extract_arguments(form)) for bv in self.get_dependencies(): c_rep = bv.saved_output tlm_value = bv.tlm_value @@ -79,6 +118,9 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar if not isinstance(dform, float): dform = ufl.algorithms.expand_derivatives(dform) dform = self.compat.assemble_adjoint_value(dform) + if arity_form == 1 and dform != 0: + # Then dform is a Vector + dform = dform.function return dform def prepare_evaluate_hessian(self, inputs, hessian_inputs, adj_inputs, relevant_dependencies): @@ -90,27 +132,27 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v hessian_input = hessian_inputs[0] adj_input = adj_inputs[0] + from ufl.algorithms.analysis import extract_arguments + arity_form = len(extract_arguments(form)) + c1 = block_variable.output c1_rep = block_variable.saved_output if isinstance(c1, self.backend.Function): - dc = self.backend.TestFunction(c1.function_space()) + space = c1.function_space() elif isinstance(c1, self.compat.ExpressionType): mesh = form.ufl_domain().ufl_cargo() - W = c1._ad_function_space(mesh) - dc = self.backend.TestFunction(W) + space = c1._ad_function_space(mesh) elif isinstance(c1, self.backend.Constant): mesh = self.compat.extract_mesh_from_form(form) - dc = self.backend.TestFunction(c1._ad_function_space(mesh)) + space = c1._ad_function_space(mesh) elif isinstance(c1, self.compat.MeshType): c1_rep = self.backend.SpatialCoordinate(c1) - dc = self.backend.TestFunction(c1._ad_function_space()) + space = c1._ad_function_space() else: return None - dform = self.backend.derivative(form, c1_rep, dc) - dform = ufl.algorithms.expand_derivatives(dform) - hessian_outputs = hessian_input * self.compat.assemble_adjoint_value(dform) + hessian_outputs, dform = self.compute_action_adjoint(hessian_input, arity_form, form, c1_rep, space) ddform = 0 for other_idx, bv in relevant_dependencies: @@ -129,10 +171,10 @@ def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_v if not isinstance(ddform, float): ddform = ufl.algorithms.expand_derivatives(ddform) if not ddform.empty(): - hessian_outputs += adj_input * self.compat.assemble_adjoint_value(ddform) + hessian_outputs += self.compute_action_adjoint(adj_input, arity_form, dform=ddform)[0] if isinstance(c1, self.compat.ExpressionType): - return [(hessian_outputs, W)] + return [(hessian_outputs, space)] else: return hessian_outputs From b9f84c0ff89872a6fa194c9d07a58de451270bd5 Mon Sep 17 00:00:00 2001 From: nbouziani Date: Wed, 19 Oct 2022 13:21:36 +0100 Subject: [PATCH 2/7] Add tests --- tests/firedrake_adjoint/test_assemble.py | 89 +++++++++++++++++++++++- 1 file changed, 86 insertions(+), 3 deletions(-) diff --git a/tests/firedrake_adjoint/test_assemble.py b/tests/firedrake_adjoint/test_assemble.py index ac53b1d6..7098808c 100644 --- a/tests/firedrake_adjoint/test_assemble.py +++ b/tests/firedrake_adjoint/test_assemble.py @@ -1,9 +1,12 @@ -import pytest -pytest.importorskip("firedrake") - from firedrake import * from firedrake_adjoint import * + from numpy.testing import assert_approx_equal +from numpy.random import rand + +import pytest +pytest.importorskip("firedrake") + def test_assemble_0_forms(): mesh = IntervalMesh(10, 0, 1) @@ -21,6 +24,7 @@ def test_assemble_0_forms(): # derivative is: (1+2*u+6*u**2)*dx - summing is equivalent to testing with 1 assert_approx_equal(rf.derivative().vector().sum(), 1. + 2. * 4 + 6 * 16.) + def test_assemble_0_forms_mixed(): mesh = IntervalMesh(10, 0, 1) V = FunctionSpace(mesh, "Lagrange", 1) @@ -37,3 +41,82 @@ def test_assemble_0_forms_mixed(): rf = ReducedFunctional(s, Control(u)) # derivative is: (1+4*u)*dx - summing is equivalent to testing with 1 assert_approx_equal(rf.derivative().vector().sum(), 1. + 4. * 7) + + +def test_assemble_1_forms_adjoint(): + mesh = IntervalMesh(10, 0, 1) + V = FunctionSpace(mesh, "Lagrange", 1) + v = TestFunction(V) + f = Function(V).assign(1) + + def J(f): + w1 = assemble(inner(f, v) * dx) + w2 = assemble(inner(f**2, v) * dx) + w3 = assemble(inner(f**3, v) * dx) + return assemble((w1 + w2 + w3)**2 * dx) + + _test_adjoint(J, f) + + +def test_assemble_1_forms_tlm(): + tape = Tape() + set_working_tape(tape) + + mesh = IntervalMesh(10, 0, 1) + V = FunctionSpace(mesh, "Lagrange", 1) + v = TestFunction(V) + f = Function(V).assign(1) + + w1 = assemble(inner(f, v) * dx) + w2 = assemble(inner(f**2, v) * dx) + w3 = assemble(inner(f**3, v) * dx) + J = assemble((w1 + w2 + w3)**2 * dx) + + Jhat = ReducedFunctional(J, Control(f)) + h = Function(V) + h.vector()[:] = rand(h.dof_dset.size) + g = f.copy(deepcopy=True) + f.block_variable.tlm_value = h + tape.evaluate_tlm() + assert (taylor_test(Jhat, g, h, dJdm=J.block_variable.tlm_value) > 1.9) + + +def _test_adjoint(J, f): + import numpy.random + tape = Tape() + set_working_tape(tape) + + V = f.function_space() + h = Function(V) + h.vector()[:] = numpy.random.rand(V.dim()) + + eps_ = [0.01 / 2.0**i for i in range(5)] + residuals = [] + for eps in eps_: + + Jp = J(f + eps * h) + tape.clear_tape() + Jm = J(f) + Jm.block_variable.adj_value = 1.0 + tape.evaluate_adj() + + dJdf = f.block_variable.adj_value + + residual = abs(Jp - Jm - eps * dJdf.inner(h.vector())) + residuals.append(residual) + + r = convergence_rates(residuals, eps_) + print(r) + print(residuals) + + tol = 1E-1 + assert(r[-1] > 2 - tol) + + +def convergence_rates(E_values, eps_values): + from numpy import log + r = [] + for i in range(1, len(eps_values)): + r.append(log(E_values[i] / E_values[i - 1]) / log(eps_values[i] / eps_values[i - 1])) + + return r From fa26e979f6c4601b7082fd5230ae7e425d9911cc Mon Sep 17 00:00:00 2001 From: nbouziani Date: Wed, 19 Oct 2022 13:21:58 +0100 Subject: [PATCH 3/7] Add firedrake branch --- .circleci/config.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.circleci/config.yml b/.circleci/config.yml index 12151ac7..a69307e7 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -96,6 +96,8 @@ jobs: - run: name: install dependencies command: | + cd /home/firedrake/firedrake/src/firedrake + git checkout annotation_assemble_1-form source /home/firedrake/firedrake/bin/activate pip3 install -e .[test] From 381e8a259fae15e1279905bd4f64211099828090 Mon Sep 17 00:00:00 2001 From: nbouziani Date: Mon, 24 Oct 2022 15:13:40 +0100 Subject: [PATCH 4/7] Update spatial derivatives case and add docstring --- dolfin_adjoint_common/blocks/assembly.py | 15 +++++++++++---- tests/firedrake_adjoint/test_assemble.py | 3 ++- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/dolfin_adjoint_common/blocks/assembly.py b/dolfin_adjoint_common/blocks/assembly.py index a8aa7360..95dda3b0 100644 --- a/dolfin_adjoint_common/blocks/assembly.py +++ b/dolfin_adjoint_common/blocks/assembly.py @@ -19,6 +19,15 @@ def __str__(self): return f"assemble({ufl2unicode(self.form)})" def compute_action_adjoint(self, adj_input, arity_form, form=None, c_rep=None, space=None, dform=None): + """This computes the action of the adjoint of the derivative of `form` wrt `c_rep` on `adj_input`. + In other words, it returns: `<(dform/dc_rep)*, adj_input>` + + - If `form` has arity 0 => `dform/dc_rep` is a 1-form and `adj_input` a foat, we can simply use + the `*` operator. + + - If `form` has arity 1 => `dform/dc_rep` is a 2-form and we can symbolically take its adjoint + and then apply the action on `adj_input`, to finally assemble the result. + """ if arity_form == 0: if dform is None: dc = self.backend.TestFunction(space) @@ -42,13 +51,11 @@ def compute_action_adjoint(self, adj_input, arity_form, form=None, c_rep=None, s else: # Get PETSc matrix dform_mat = self.compat.assemble_adjoint_value(dform).petscmat - # Adjoint (hermitian transpose) - dform_mat.hermitianTranspose() - # Action + # Action of the adjoint (Hermitian transpose) adj_output = self.backend.Function(space) with adj_input.dat.vec_ro as v_vec: with adj_output.dat.vec as res_vec: - dform_mat.mult(v_vec, res_vec) + dform_mat.multHermitian(v_vec, res_vec) return adj_output, dform else: raise ValueError('Forms with arity > 1 are not handled yet!') diff --git a/tests/firedrake_adjoint/test_assemble.py b/tests/firedrake_adjoint/test_assemble.py index 7098808c..9c10eca8 100644 --- a/tests/firedrake_adjoint/test_assemble.py +++ b/tests/firedrake_adjoint/test_assemble.py @@ -47,7 +47,8 @@ def test_assemble_1_forms_adjoint(): mesh = IntervalMesh(10, 0, 1) V = FunctionSpace(mesh, "Lagrange", 1) v = TestFunction(V) - f = Function(V).assign(1) + x, = SpatialCoordinate(mesh) + f = Function(V).interpolate(cos(x)) def J(f): w1 = assemble(inner(f, v) * dx) From ce811ae242aadbef002229823be9ce82589838ab Mon Sep 17 00:00:00 2001 From: nbouziani Date: Mon, 24 Oct 2022 15:51:02 +0100 Subject: [PATCH 5/7] Fix lint --- dolfin_adjoint_common/blocks/function.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dolfin_adjoint_common/blocks/function.py b/dolfin_adjoint_common/blocks/function.py index 23b1d767..9dc200d8 100644 --- a/dolfin_adjoint_common/blocks/function.py +++ b/dolfin_adjoint_common/blocks/function.py @@ -13,7 +13,7 @@ def __init__(self, func, other, ad_block_tag=None): other = AdjFloat(other) if isinstance(other, OverloadedType): self.add_dependency(other, no_duplicates=True) - elif not(isinstance(other, float) or isinstance(other, int)): + elif not (isinstance(other, float) or isinstance(other, int)): # Assume that this is a point-wise evaluated UFL expression (firedrake only) for op in traverse_unique_terminals(other): if isinstance(op, OverloadedType): From 6994441f651a01048dd14118b591f35f6712f49b Mon Sep 17 00:00:00 2001 From: nbouziani Date: Sat, 19 Nov 2022 14:58:57 +0000 Subject: [PATCH 6/7] Swap adjoint action for arity-0 case for type compatibility (always need to return a Vector) --- dolfin_adjoint_common/blocks/assembly.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/dolfin_adjoint_common/blocks/assembly.py b/dolfin_adjoint_common/blocks/assembly.py index 95dda3b0..af5a9ea1 100644 --- a/dolfin_adjoint_common/blocks/assembly.py +++ b/dolfin_adjoint_common/blocks/assembly.py @@ -33,7 +33,8 @@ def compute_action_adjoint(self, adj_input, arity_form, form=None, c_rep=None, s dc = self.backend.TestFunction(space) dform = self.backend.derivative(form, c_rep, dc) dform_vector = self.compat.assemble_adjoint_value(dform) - return adj_input * dform_vector, dform + # Return a Vector scaled by the scalar `adj_input` + return dform_vector * adj_input, dform elif arity_form == 1: if dform is None: dc = self.backend.TrialFunction(space) From 9f9b0ccc0fcb5cf7d1375f3de013e26bfd750a65 Mon Sep 17 00:00:00 2001 From: nbouziani Date: Sat, 19 Nov 2022 15:01:56 +0000 Subject: [PATCH 7/7] Update config.yml as firedrake PR got merged --- .circleci/config.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index a69307e7..12151ac7 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -96,8 +96,6 @@ jobs: - run: name: install dependencies command: | - cd /home/firedrake/firedrake/src/firedrake - git checkout annotation_assemble_1-form source /home/firedrake/firedrake/bin/activate pip3 install -e .[test]