From 2880f2be69747385ea77eb42296ede967ee15c86 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Wed, 10 Jul 2024 10:55:43 +0100 Subject: [PATCH 01/28] Reverse-over-forward configuration controls --- pyadjoint/__init__.py | 5 ++- pyadjoint/tape.py | 62 ++++++++++++++++++++++++++++++++++++ tests/pyadjoint/test_tape.py | 29 +++++++++++++++++ 3 files changed, 95 insertions(+), 1 deletion(-) create mode 100644 tests/pyadjoint/test_tape.py diff --git a/pyadjoint/__init__.py b/pyadjoint/__init__.py index 0baeb05b..b923bfa7 100644 --- a/pyadjoint/__init__.py +++ b/pyadjoint/__init__.py @@ -10,7 +10,10 @@ from .block import Block from .tape import (Tape, set_working_tape, get_working_tape, no_annotations, - annotate_tape, stop_annotating, pause_annotation, continue_annotation) + annotate_tape, stop_annotating, pause_annotation, continue_annotation, + no_reverse_over_forward, reverse_over_forward_enabled, + stop_reverse_over_forward, pause_reverse_over_forward, + continue_reverse_over_forward) from .adjfloat import AdjFloat, exp, log from .reduced_functional import ReducedFunctional from .drivers import compute_gradient, compute_hessian, solve_adjoint diff --git a/pyadjoint/tape.py b/pyadjoint/tape.py index 63853d44..b1668ded 100644 --- a/pyadjoint/tape.py +++ b/pyadjoint/tape.py @@ -11,6 +11,7 @@ _working_tape = None _annotation_enabled = False +_reverse_over_forward_enabled = False def get_working_tape(): @@ -136,6 +137,67 @@ def annotate_tape(kwargs=None): return annotate +def pause_reverse_over_forward(): + """Disable reverse-over-forward AD. + """ + + global _reverse_over_forward_enabled + _reverse_over_forward_enabled = False + + +def continue_reverse_over_forward(): + """Enable reverse-over-forward AD. + + Returns: + bool: True + """ + + global _reverse_over_forward_enabled + _reverse_over_forward_enabled = True + # Following continue_annotation behavior + return _reverse_over_forward_enabled + + +@contextmanager +def stop_reverse_over_forward(): + """Return a callable used to construct a context manager within which + reverse-over-forward AD is disabled. + + Returns: + callable: Callable which returns a context manager. + """ + + global _reverse_over_forward_enabled + reverse_over_forward_enabled = _reverse_over_forward_enabled + _reverse_over_forward_enabled = False + try: + yield + finally: + _reverse_over_forward_enabled = reverse_over_forward_enabled + + +def no_reverse_over_forward(function): + """Decorator to disable reverse-over-forward AD for the decorated callable. + + Args: + function (callable): The callable. + Returns: + callable: Callable for which reverse-over-forward AD is disabled. + """ + + return stop_reverse_over_forward()(function) + + +def reverse_over_forward_enabled(): + """Return whether reverse-over-forward AD is enabled. + + Returns: + bool: Whether reverse-over-forward AD is enabled. + """ + + return _reverse_over_forward_enabled + + def _find_relevant_nodes(tape, controls): # This function is just a stripped down Block.optimize_for_controls blocks = tape.get_blocks() diff --git a/tests/pyadjoint/test_tape.py b/tests/pyadjoint/test_tape.py new file mode 100644 index 00000000..580132bb --- /dev/null +++ b/tests/pyadjoint/test_tape.py @@ -0,0 +1,29 @@ +from pyadjoint import * # noqa: F403 + + +def test_reverse_over_forward_configuration(): + assert not reverse_over_forward_enabled() + + continue_reverse_over_forward() + assert reverse_over_forward_enabled() + pause_reverse_over_forward() + assert not reverse_over_forward_enabled() + + continue_reverse_over_forward() + assert reverse_over_forward_enabled() + with stop_reverse_over_forward(): + assert not reverse_over_forward_enabled() + assert reverse_over_forward_enabled() + pause_reverse_over_forward() + assert not reverse_over_forward_enabled() + + @no_reverse_over_forward + def test(): + assert not reverse_over_forward_enabled() + + continue_reverse_over_forward() + assert reverse_over_forward_enabled() + test() + assert reverse_over_forward_enabled() + pause_reverse_over_forward() + assert not reverse_over_forward_enabled() From 19e7a1d5bf89108014d91030aa307d3995c1b2be Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Wed, 10 Jul 2024 11:06:33 +0100 Subject: [PATCH 02/28] Reverse-over-forward AD --- pyadjoint/block.py | 27 +++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/pyadjoint/block.py b/pyadjoint/block.py index 29dbc70b..16bf7f2e 100644 --- a/pyadjoint/block.py +++ b/pyadjoint/block.py @@ -1,4 +1,4 @@ -from .tape import no_annotations +from .tape import no_annotations, reverse_over_forward_enabled from html import escape @@ -11,15 +11,19 @@ class Block(object): Abstract methods :func:`evaluate_adj` + Args: + n_outputs (int): The number of outputs. Required for + reverse-over-forward AD. """ __slots__ = ['_dependencies', '_outputs', 'block_helper'] pop_kwargs_keys = [] - def __init__(self, ad_block_tag=None): + def __init__(self, ad_block_tag=None, *, n_outputs=1): self._dependencies = [] self._outputs = [] self.block_helper = None self.tag = ad_block_tag + self._n_outputs = n_outputs @classmethod def pop_kwargs(cls, kwargs): @@ -71,9 +75,19 @@ def add_output(self, obj): obj (:class:`BlockVariable`): The object to be added. """ + + if reverse_over_forward_enabled() and len(self._outputs) >= self._n_outputs: + raise RuntimeError("Unexpected output") + obj.will_add_as_output() self._outputs.append(obj) + if reverse_over_forward_enabled(): + if len(self._outputs) == self._n_outputs: + self.solve_tlm() + elif len(self._outputs) > self._n_outputs: + raise RuntimeError("Unexpected output") + def get_outputs(self): """Returns the list of block outputs. @@ -255,6 +269,15 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar """ raise NotImplementedError("evaluate_tlm_component is not implemented for Block-type: {}".format(type(self))) + def solve_tlm(self): + """This method should be overridden if using reverse-over-forward AD. + + Perform a tangent-linear operation, storing results in the `tlm_value` + attributes of relevant `BlockVariable` objects. + """ + + raise NotImplementedError(f"solve_tlm is not implemented for Block-type: {type(self)}") + @no_annotations def evaluate_hessian(self, markings=False): outputs = self.get_outputs() From 857438f7a126704af000d2963b467c526cc111c5 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Wed, 10 Jul 2024 11:25:26 +0100 Subject: [PATCH 03/28] Test setup --- tests/pyadjoint/test_tape.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/pyadjoint/test_tape.py b/tests/pyadjoint/test_tape.py index 580132bb..c81637eb 100644 --- a/tests/pyadjoint/test_tape.py +++ b/tests/pyadjoint/test_tape.py @@ -1,6 +1,15 @@ +import pytest + from pyadjoint import * # noqa: F403 +@pytest.fixture(autouse=True, scope="module") +def _(): + pause_reverse_over_forward() + yield + pause_reverse_over_forward() + + def test_reverse_over_forward_configuration(): assert not reverse_over_forward_enabled() From 04c48d622680b7a98c5775209f07a159e81ba695 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Wed, 10 Jul 2024 13:06:40 +0100 Subject: [PATCH 04/28] In-place assignment, restore old values in reverse-over-forward AD --- pyadjoint/block.py | 9 +++++++-- pyadjoint/block_variable.py | 21 +++++++++++++++++++++ pyadjoint/overloaded_type.py | 12 ++++++++++++ 3 files changed, 40 insertions(+), 2 deletions(-) diff --git a/pyadjoint/block.py b/pyadjoint/block.py index 16bf7f2e..52ee4d85 100644 --- a/pyadjoint/block.py +++ b/pyadjoint/block.py @@ -1,6 +1,8 @@ -from .tape import no_annotations, reverse_over_forward_enabled +from contextlib import ExitStack from html import escape +from .tape import no_annotations, reverse_over_forward_enabled + class Block(object): """Base class for all Tape Block types. @@ -84,7 +86,10 @@ def add_output(self, obj): if reverse_over_forward_enabled(): if len(self._outputs) == self._n_outputs: - self.solve_tlm() + with ExitStack() as stack: + for output in self._outputs: + stack.enter_context(output.restore_output()) + self.solve_tlm() elif len(self._outputs) > self._n_outputs: raise RuntimeError("Unexpected output") diff --git a/pyadjoint/block_variable.py b/pyadjoint/block_variable.py index e252b0ea..1fd70a5a 100644 --- a/pyadjoint/block_variable.py +++ b/pyadjoint/block_variable.py @@ -1,3 +1,5 @@ +from contextlib import contextmanager + from .tape import no_annotations, get_working_tape @@ -93,3 +95,22 @@ def checkpoint(self, value): if self.is_control: return self._checkpoint = value + + @contextmanager + def restore_output(self): + """Return a context manager which can be used to temporarily restore the + value of `self.output` to `self.block_variable.saved_output`. + + Returns: + The context manager + """ + + if self.output is self.saved_output: + yield + else: + old_value = self.output._ad_copy() + self.output._ad_assign(self.saved_output) + try: + yield + finally: + self._output._ad_assign(old_value) diff --git a/pyadjoint/overloaded_type.py b/pyadjoint/overloaded_type.py index 0f96667e..d2f9b183 100644 --- a/pyadjoint/overloaded_type.py +++ b/pyadjoint/overloaded_type.py @@ -285,6 +285,18 @@ def _ad_to_list(m): """ raise NotImplementedError + def _ad_assign(self, other): + """This method must be overridden for mutable types. + + In-place assignment. + + Args: + other (object): The object assign to `self`, with the same type as + `self`. + """ + + raise NotImplementedError + def _ad_copy(self): """This method must be overridden. From 27d6355ed742d23fcef06595005d829f6b880fb5 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Wed, 10 Jul 2024 15:53:17 +0100 Subject: [PATCH 05/28] Test setup fix --- tests/pyadjoint/test_tape.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pyadjoint/test_tape.py b/tests/pyadjoint/test_tape.py index c81637eb..20b88189 100644 --- a/tests/pyadjoint/test_tape.py +++ b/tests/pyadjoint/test_tape.py @@ -3,7 +3,7 @@ from pyadjoint import * # noqa: F403 -@pytest.fixture(autouse=True, scope="module") +@pytest.fixture(autouse=True) def _(): pause_reverse_over_forward() yield From d123af3c19139631c78f9742bca02eb370750923 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Wed, 10 Jul 2024 17:01:33 +0100 Subject: [PATCH 06/28] BlockVariable.restore_output fix --- pyadjoint/block_variable.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/pyadjoint/block_variable.py b/pyadjoint/block_variable.py index 1fd70a5a..c59d8f19 100644 --- a/pyadjoint/block_variable.py +++ b/pyadjoint/block_variable.py @@ -1,6 +1,6 @@ from contextlib import contextmanager -from .tape import no_annotations, get_working_tape +from .tape import no_annotations, get_working_tape, stop_annotating class BlockVariable(object): @@ -108,9 +108,10 @@ def restore_output(self): if self.output is self.saved_output: yield else: - old_value = self.output._ad_copy() - self.output._ad_assign(self.saved_output) - try: - yield - finally: - self._output._ad_assign(old_value) + with stop_annotating(): + old_value = self.output._ad_copy() + self.output._ad_assign(self.saved_output) + try: + yield + finally: + self.output._ad_assign(old_value) From ad76a174a397e4f3adda0d0062d62af36eb45b54 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Wed, 10 Jul 2024 18:33:58 +0100 Subject: [PATCH 07/28] Move will_add_as_output call until after tangent-linear operations --- pyadjoint/block.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyadjoint/block.py b/pyadjoint/block.py index 52ee4d85..a20a304d 100644 --- a/pyadjoint/block.py +++ b/pyadjoint/block.py @@ -81,7 +81,6 @@ def add_output(self, obj): if reverse_over_forward_enabled() and len(self._outputs) >= self._n_outputs: raise RuntimeError("Unexpected output") - obj.will_add_as_output() self._outputs.append(obj) if reverse_over_forward_enabled(): @@ -93,6 +92,8 @@ def add_output(self, obj): elif len(self._outputs) > self._n_outputs: raise RuntimeError("Unexpected output") + obj.will_add_as_output() + def get_outputs(self): """Returns the list of block outputs. From 8b0d18c918c45ee3d826e5365b25d8a1d0f29350 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Wed, 10 Jul 2024 19:47:30 +0100 Subject: [PATCH 08/28] Handle zero case to avoid unnecessary higher-order processing --- pyadjoint/block.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pyadjoint/block.py b/pyadjoint/block.py index a20a304d..9e98b6a7 100644 --- a/pyadjoint/block.py +++ b/pyadjoint/block.py @@ -85,10 +85,14 @@ def add_output(self, obj): if reverse_over_forward_enabled(): if len(self._outputs) == self._n_outputs: - with ExitStack() as stack: - for output in self._outputs: - stack.enter_context(output.restore_output()) - self.solve_tlm() + if any(dep.tlm_value is not None for dep in self.get_dependencies()): + with ExitStack() as stack: + for output in self._outputs: + stack.enter_context(output.restore_output()) + self.solve_tlm() + else: + for x in self.get_outputs(): + x.tlm_value = None elif len(self._outputs) > self._n_outputs: raise RuntimeError("Unexpected output") From 111fb9feb6347b55cb131ae3048a56204a34c049 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 13:16:58 +0100 Subject: [PATCH 09/28] restore_output fixes --- pyadjoint/block.py | 4 ++-- pyadjoint/block_variable.py | 11 ++++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/pyadjoint/block.py b/pyadjoint/block.py index 9e98b6a7..caf6ef2d 100644 --- a/pyadjoint/block.py +++ b/pyadjoint/block.py @@ -87,8 +87,8 @@ def add_output(self, obj): if len(self._outputs) == self._n_outputs: if any(dep.tlm_value is not None for dep in self.get_dependencies()): with ExitStack() as stack: - for output in self._outputs: - stack.enter_context(output.restore_output()) + for dep in self.get_dependencies(): + stack.enter_context(dep.restore_output()) self.solve_tlm() else: for x in self.get_outputs(): diff --git a/pyadjoint/block_variable.py b/pyadjoint/block_variable.py index c59d8f19..9c4ae10c 100644 --- a/pyadjoint/block_variable.py +++ b/pyadjoint/block_variable.py @@ -98,8 +98,8 @@ def checkpoint(self, value): @contextmanager def restore_output(self): - """Return a context manager which can be used to temporarily restore the - value of `self.output` to `self.block_variable.saved_output`. + """Return a context manager which can be used to temporarily restore + the value of `self.output` to `self.block_variable.saved_output`. Returns: The context manager @@ -111,7 +111,8 @@ def restore_output(self): with stop_annotating(): old_value = self.output._ad_copy() self.output._ad_assign(self.saved_output) - try: - yield - finally: + try: + yield + finally: + with stop_annotating(): self.output._ad_assign(old_value) From 74e45feea90a0951e60c181eb68fe4dc679459ab Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 13:44:59 +0100 Subject: [PATCH 10/28] Limit reverse-over-forward to second order --- pyadjoint/block.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pyadjoint/block.py b/pyadjoint/block.py index caf6ef2d..9d8ebdc2 100644 --- a/pyadjoint/block.py +++ b/pyadjoint/block.py @@ -1,7 +1,7 @@ from contextlib import ExitStack from html import escape -from .tape import no_annotations, reverse_over_forward_enabled +from .tape import no_annotations, reverse_over_forward_enabled, stop_reverse_over_forward class Block(object): @@ -89,7 +89,8 @@ def add_output(self, obj): with ExitStack() as stack: for dep in self.get_dependencies(): stack.enter_context(dep.restore_output()) - self.solve_tlm() + with stop_reverse_over_forward(): + self.solve_tlm() else: for x in self.get_outputs(): x.tlm_value = None From adcad2ddb6844cfe5ae9f816b278b1d69353482c Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 14:06:50 +0100 Subject: [PATCH 11/28] Reverse-over-forward AD: ExpBlock --- pyadjoint/adjfloat.py | 8 ++++++ tests/pyadjoint/test_reverse_over_forward.py | 27 ++++++++++++++++++++ 2 files changed, 35 insertions(+) create mode 100644 tests/pyadjoint/test_reverse_over_forward.py diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index e92e2a05..67d10384 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -188,6 +188,14 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar input0 = inputs[0] return _exp(input0) * tlm_input + def solve_tlm(self): + x, = self.get_outputs() + a, = self.get_dependencies() + if a.tlm_value is None: + x.tlm_value = None + else: + x.tlm_value = exp(a.output) * a.tlm_value + def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): input0 = inputs[0] diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py new file mode 100644 index 00000000..eeee04c0 --- /dev/null +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -0,0 +1,27 @@ +from contextlib import contextmanager + +import pytest + +from pyadjoint import * + + +@pytest.fixture(autouse=True) +def _(): + get_working_tape().clear_tape() + continue_annotation() + continue_reverse_over_forward() + yield + get_working_tape().clear_tape() + pause_annotation() + pause_reverse_over_forward() + + +@pytest.mark.parametrize("a_val", [2.0, -2.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +def test_exp(a_val, tlm_a_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = exp(a) + _ = compute_gradient(b.block_variable.tlm_value, Control(a)) + adj_value = a.block_variable.adj_value + assert adj_value == exp(a_val) * tlm_a_val From 3724a74b7b1b6d3939bef9607ee876a189e9f312 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 14:10:02 +0100 Subject: [PATCH 12/28] Reverse-over-forward AD: LogBlock --- pyadjoint/adjfloat.py | 8 ++++++++ tests/pyadjoint/test_reverse_over_forward.py | 11 +++++++++++ 2 files changed, 19 insertions(+) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 67d10384..34bb7127 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -221,6 +221,14 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar input0 = inputs[0] return tlm_input / input0 + def solve_tlm(self): + x, = self.get_outputs() + a, = self.get_dependencies() + if a.tlm_value is None: + x.tlm_value = None + else: + x.tlm_value = a.tlm_value / a.output + def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): input0 = inputs[0] diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index eeee04c0..799fd440 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -25,3 +25,14 @@ def test_exp(a_val, tlm_a_val): _ = compute_gradient(b.block_variable.tlm_value, Control(a)) adj_value = a.block_variable.adj_value assert adj_value == exp(a_val) * tlm_a_val + + +@pytest.mark.parametrize("a_val", [2.0, 3.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +def test_log(a_val, tlm_a_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = log(a) + _ = compute_gradient(b.block_variable.tlm_value, Control(a)) + adj_value = a.block_variable.adj_value + assert adj_value == -tlm_a_val / (a_val ** 2) From 7e77ccdbd560605ff5ebe55fe6d4f7a767045cbd Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 14:23:25 +0100 Subject: [PATCH 13/28] Reverse-over-forward AD: AddBlock --- pyadjoint/adjfloat.py | 11 ++++++++++ tests/pyadjoint/test_reverse_over_forward.py | 21 ++++++++++++++++---- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 34bb7127..14d6599d 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -458,6 +458,17 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar tlm_output += tlm_input return tlm_output + def solve_tlm(self): + x, = self.get_outputs() + terms = tuple(dep.tlm_value for dep in self.get_dependencies() + if dep.tlm_value is not None) + if len(terms) == 0: + x.tlm_value = None + elif len(terms) == 1: + x.tlm_value = AdjFloat(terms[0]) + else: + x.tlm_value = sum(terms[1:], start=terms[0]) + def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): return hessian_inputs[0] diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index 799fd440..e8ced07f 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -21,8 +21,8 @@ def _(): def test_exp(a_val, tlm_a_val): a = AdjFloat(a_val) a.block_variable.tlm_value = AdjFloat(tlm_a_val) - b = exp(a) - _ = compute_gradient(b.block_variable.tlm_value, Control(a)) + x = exp(a) + _ = compute_gradient(x.block_variable.tlm_value, Control(a)) adj_value = a.block_variable.adj_value assert adj_value == exp(a_val) * tlm_a_val @@ -32,7 +32,20 @@ def test_exp(a_val, tlm_a_val): def test_log(a_val, tlm_a_val): a = AdjFloat(a_val) a.block_variable.tlm_value = AdjFloat(tlm_a_val) - b = log(a) - _ = compute_gradient(b.block_variable.tlm_value, Control(a)) + x = log(a) + _ = compute_gradient(x.block_variable.tlm_value, Control(a)) adj_value = a.block_variable.adj_value assert adj_value == -tlm_a_val / (a_val ** 2) + + +@pytest.mark.parametrize("a_val", [2.0, -2.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("b_val", [4.25, -4.25]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125]) +def test_add(a_val, tlm_a_val, b_val, tlm_b_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = AdjFloat(b_val) + b.block_variable.tlm_value = AdjFloat(tlm_b_val) + x = a + b + assert x.block_variable.tlm_value == tlm_a_val + tlm_b_val From c3ad8ea6dec51563ff5a3ce9a5780f47c92fdf93 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 14:27:35 +0100 Subject: [PATCH 14/28] Reverse-over-forward AD: NegBlock --- pyadjoint/adjfloat.py | 8 ++++++++ tests/pyadjoint/test_reverse_over_forward.py | 9 +++++++++ 2 files changed, 17 insertions(+) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 14d6599d..ff9ef775 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -625,6 +625,14 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None): return float.__neg__(tlm_inputs[0]) + def solve_tlm(self): + x, = self.get_outputs() + a, = self.get_dependencies() + if a.tlm_value is None: + x.tlm_value = None + else: + x.tlm_value = -a.tlm_value + def evaluate_hessian(self, markings=False): hessian_input = self.get_outputs()[0].hessian_value if hessian_input is None: diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index e8ced07f..fce60421 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -49,3 +49,12 @@ def test_add(a_val, tlm_a_val, b_val, tlm_b_val): b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = a + b assert x.block_variable.tlm_value == tlm_a_val + tlm_b_val + + +@pytest.mark.parametrize("a_val", [2.0, -2.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +def test_neg(a_val, tlm_a_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + x = -a + assert x.block_variable.tlm_value == -tlm_a_val From f585e1886a994cbf8e87ebb0b0ba5226fefdfd40 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 14:30:32 +0100 Subject: [PATCH 15/28] Reverse-over-forward AD: SubBlock --- pyadjoint/adjfloat.py | 13 +++++++++++++ tests/pyadjoint/test_reverse_over_forward.py | 13 +++++++++++++ 2 files changed, 26 insertions(+) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index ff9ef775..7d829482 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -493,6 +493,19 @@ def evaluate_tlm(self, markings=False): if tlm_input_1 is not None: output.add_tlm_output(float.__neg__(tlm_input_1)) + def solve_tlm(self): + x, = self.get_outputs() + a, b = self.get_dependencies() + if a.tlm_value is None: + if b.tlm_value is None: + x.tlm_value = None + else: + x.tlm_value = -b.tlm_value + elif b.tlm_value is None: + x.tlm_value = a.tlm_value + else: + x.tlm_value = a.tlm_value - b.tlm_value + def evaluate_hessian(self, markings=False): hessian_input = self.get_outputs()[0].hessian_value if hessian_input is None: diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index fce60421..cb10efdc 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -51,6 +51,19 @@ def test_add(a_val, tlm_a_val, b_val, tlm_b_val): assert x.block_variable.tlm_value == tlm_a_val + tlm_b_val +@pytest.mark.parametrize("a_val", [2.0, -2.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("b_val", [4.25, -4.25]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125]) +def test_sub(a_val, tlm_a_val, b_val, tlm_b_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = AdjFloat(b_val) + b.block_variable.tlm_value = AdjFloat(tlm_b_val) + x = a - b + assert x.block_variable.tlm_value == tlm_a_val - tlm_b_val + + @pytest.mark.parametrize("a_val", [2.0, -2.0]) @pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) def test_neg(a_val, tlm_a_val): From f537007ab008758f8a6acbce252c9ff82bdeaaf6 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 14:33:48 +0100 Subject: [PATCH 16/28] Expand AddBlock and SubBlock reverse-over-forward tests --- tests/pyadjoint/test_reverse_over_forward.py | 32 ++++++++++++++------ 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index cb10efdc..5b3b2272 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -39,29 +39,41 @@ def test_log(a_val, tlm_a_val): @pytest.mark.parametrize("a_val", [2.0, -2.0]) -@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5, None]) @pytest.mark.parametrize("b_val", [4.25, -4.25]) -@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125, None]) def test_add(a_val, tlm_a_val, b_val, tlm_b_val): a = AdjFloat(a_val) - a.block_variable.tlm_value = AdjFloat(tlm_a_val) + if tlm_a_val is not None: + a.block_variable.tlm_value = AdjFloat(tlm_a_val) b = AdjFloat(b_val) - b.block_variable.tlm_value = AdjFloat(tlm_b_val) + if tlm_b_val is not None: + b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = a + b - assert x.block_variable.tlm_value == tlm_a_val + tlm_b_val + if tlm_a_val is None and tlm_b_val is None: + assert x.block_variable.tlm_value is None + else: + assert (x.block_variable.tlm_value == + (0.0 if tlm_a_val is None else tlm_a_val) + (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) -@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5, None]) @pytest.mark.parametrize("b_val", [4.25, -4.25]) -@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125, None]) def test_sub(a_val, tlm_a_val, b_val, tlm_b_val): a = AdjFloat(a_val) - a.block_variable.tlm_value = AdjFloat(tlm_a_val) + if tlm_a_val is not None: + a.block_variable.tlm_value = AdjFloat(tlm_a_val) b = AdjFloat(b_val) - b.block_variable.tlm_value = AdjFloat(tlm_b_val) + if tlm_b_val is not None: + b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = a - b - assert x.block_variable.tlm_value == tlm_a_val - tlm_b_val + if tlm_a_val is None and tlm_b_val is None: + assert x.block_variable.tlm_value is None + else: + assert (x.block_variable.tlm_value == + (0.0 if tlm_a_val is None else tlm_a_val) - (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) From ec60d22393aa5ecbea7b9badeaa21fd3ea06365e Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 14:43:16 +0100 Subject: [PATCH 17/28] Reverse-over-forward AD: MulBlock --- pyadjoint/adjfloat.py | 26 +++++++++++++++----- tests/pyadjoint/test_reverse_over_forward.py | 26 ++++++++++++++++++-- 2 files changed, 44 insertions(+), 8 deletions(-) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 7d829482..8ca0a095 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -440,6 +440,15 @@ def evaluate_hessian(self, markings=False): exponent.add_hessian_output(float.__mul__(base.tlm_value, mixed)) +def sum_tlm_terms(terms): + if len(terms) == 0: + return None + elif len(terms) == 1: + return AdjFloat(terms[0]) + else: + return sum(terms[1:], start=terms[0]) + + class AddBlock(FloatOperatorBlock): operator = staticmethod(float.__add__) symbol = "+" @@ -462,12 +471,7 @@ def solve_tlm(self): x, = self.get_outputs() terms = tuple(dep.tlm_value for dep in self.get_dependencies() if dep.tlm_value is not None) - if len(terms) == 0: - x.tlm_value = None - elif len(terms) == 1: - x.tlm_value = AdjFloat(terms[0]) - else: - x.tlm_value = sum(terms[1:], start=terms[0]) + x.tlm_value = sum_tlm_terms(terms) def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): @@ -534,6 +538,16 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar tlm_output += float.__mul__(tlm_input, self.terms[j].saved_output) return tlm_output + def solve_tlm(self): + x, = self.get_outputs() + a, b = self.get_dependencies() + terms = [] + if a.tlm_value is not None: + terms.append(b.output * a.tlm_value) + if b.tlm_value is not None: + terms.append(a.output * b.tlm_value) + x.tlm_value = sum_tlm_terms(terms) + def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): adj_input = adj_inputs[0] diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index 5b3b2272..3a543400 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -54,7 +54,8 @@ def test_add(a_val, tlm_a_val, b_val, tlm_b_val): assert x.block_variable.tlm_value is None else: assert (x.block_variable.tlm_value == - (0.0 if tlm_a_val is None else tlm_a_val) + (0.0 if tlm_b_val is None else tlm_b_val)) + (0.0 if tlm_a_val is None else tlm_a_val) + + (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) @@ -73,7 +74,28 @@ def test_sub(a_val, tlm_a_val, b_val, tlm_b_val): assert x.block_variable.tlm_value is None else: assert (x.block_variable.tlm_value == - (0.0 if tlm_a_val is None else tlm_a_val) - (0.0 if tlm_b_val is None else tlm_b_val)) + (0.0 if tlm_a_val is None else tlm_a_val) + - (0.0 if tlm_b_val is None else tlm_b_val)) + + +@pytest.mark.parametrize("a_val", [2.0, -2.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5, None]) +@pytest.mark.parametrize("b_val", [4.25, -4.25]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125, None]) +def test_mul(a_val, tlm_a_val, b_val, tlm_b_val): + a = AdjFloat(a_val) + if tlm_a_val is not None: + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = AdjFloat(b_val) + if tlm_b_val is not None: + b.block_variable.tlm_value = AdjFloat(tlm_b_val) + x = a * b + if tlm_a_val is None and tlm_b_val is None: + assert x.block_variable.tlm_value is None + else: + assert (x.block_variable.tlm_value == + b_val * (0.0 if tlm_a_val is None else tlm_a_val) + + a_val * (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) From 664b320475e195ec2c2aa3bf97e8734ddd69df0f Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 15:14:33 +0100 Subject: [PATCH 18/28] Reverse-over-forward AD: PowBlock --- pyadjoint/adjfloat.py | 28 +++++++++++++------ tests/pyadjoint/test_reverse_over_forward.py | 29 ++++++++++++++++++++ 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 8ca0a095..176bdf1d 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -358,6 +358,15 @@ def __str__(self): return f"{self.terms[0]} {self.symbol} {self.terms[1]}" +def sum_tlm_terms(terms): + if len(terms) == 0: + return None + elif len(terms) == 1: + return AdjFloat(terms[0]) + else: + return sum(terms[1:], start=terms[0]) + + class PowBlock(FloatOperatorBlock): operator = staticmethod(float.__pow__) symbol = "**" @@ -395,6 +404,16 @@ def evaluate_tlm(self, markings=False): float.__pow__(base_value, exponent_value)) output.add_tlm_output(exponent_adj) + def solve_tlm(self): + x, = self.get_outputs() + a, b = self.get_dependencies() + terms = [] + if a.tlm_value is not None: + terms.append(b.output * (a.output ** (b.output - 1)) * a.tlm_value) + if b.tlm_value is not None: + terms.append(log(a.output) * (a.output ** b.output) * a.tlm_value) + x.tlm_value = sum_tlm_terms(terms) + def evaluate_hessian(self, markings=False): output = self.get_outputs()[0] hessian_input = output.hessian_value @@ -440,15 +459,6 @@ def evaluate_hessian(self, markings=False): exponent.add_hessian_output(float.__mul__(base.tlm_value, mixed)) -def sum_tlm_terms(terms): - if len(terms) == 0: - return None - elif len(terms) == 1: - return AdjFloat(terms[0]) - else: - return sum(terms[1:], start=terms[0]) - - class AddBlock(FloatOperatorBlock): operator = staticmethod(float.__add__) symbol = "+" diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index 3a543400..2d87c03f 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -1,5 +1,6 @@ from contextlib import contextmanager +import numpy as np import pytest from pyadjoint import * @@ -38,6 +39,34 @@ def test_log(a_val, tlm_a_val): assert adj_value == -tlm_a_val / (a_val ** 2) +@pytest.mark.parametrize("a_val", [2.0, 3.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("b_val", [4.25, 5.25]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125]) +def test_pow(a_val, tlm_a_val, b_val, tlm_b_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = AdjFloat(b_val) + b.block_variable.tlm_value = AdjFloat(tlm_b_val) + x = a ** b + J_hat = ReducedFunctional(x.block_variable.tlm_value, Control(a)) + assert taylor_test(J_hat, a, AdjFloat(1.0)) > 1.9 + J_hat = ReducedFunctional(x.block_variable.tlm_value, Control(b)) + assert taylor_test(J_hat, b, AdjFloat(1.0)) > 1.9 + + +@pytest.mark.parametrize("a_val", [2.0, 3.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +def test_a_pow_a(a_val, tlm_a_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + x = a ** a + _ = compute_gradient(x.block_variable.tlm_value, Control(a)) + adj_value = a.block_variable.adj_value + assert np.allclose( + adj_value, ((1 + log(a_val)) ** 2 + (1 / a_val)) * (a_val ** a_val) * tlm_a_val) + + @pytest.mark.parametrize("a_val", [2.0, -2.0]) @pytest.mark.parametrize("tlm_a_val", [3.5, -3.5, None]) @pytest.mark.parametrize("b_val", [4.25, -4.25]) From 1ad632355cbfa9f841ece4d4a93962809112f24e Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 17:31:38 +0100 Subject: [PATCH 19/28] Add PosBlock, use to fix a bug in AdjFloat reverse-over-forward AD --- pyadjoint/adjfloat.py | 35 +++++++++++++++++++- tests/pyadjoint/test_reverse_over_forward.py | 9 +++++ 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 176bdf1d..d847609b 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -64,6 +64,10 @@ def __div__(self, other): def __truediv__(self, other): return DivBlock(self, other) + @annotate_operator + def __pos__(self): + return PosBlock(self) + @annotate_operator def __neg__(self): return NegBlock(self) @@ -362,7 +366,7 @@ def sum_tlm_terms(terms): if len(terms) == 0: return None elif len(terms) == 1: - return AdjFloat(terms[0]) + return +terms[0] else: return sum(terms[1:], start=terms[0]) @@ -652,6 +656,35 @@ def evaluate_hessian(self, markings=False): denominator.add_hessian_output(float.__mul__(numerator.tlm_value, mixed)) +class PosBlock(FloatOperatorBlock): + operator = staticmethod(float.__pos__) + symbol = "+" + + def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepared=None): + return float.__pos__(adj_inputs[0]) + + def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None): + return float.__pos__(tlm_inputs[0]) + + def solve_tlm(self): + x, = self.get_outputs() + a, = self.get_dependencies() + if a.tlm_value is None: + x.tlm_value = None + else: + x.tlm_value = +a.tlm_value + + def evaluate_hessian(self, markings=False): + hessian_input = self.get_outputs()[0].hessian_value + if hessian_input is None: + return + + self.terms[0].add_hessian_output(float.__pos__(hessian_input)) + + def __str__(self): + return f"{self.symbol} {self.terms[0]}" + + class NegBlock(FloatOperatorBlock): operator = staticmethod(float.__neg__) symbol = "-" diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index 2d87c03f..15ec7c25 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -127,6 +127,15 @@ def test_mul(a_val, tlm_a_val, b_val, tlm_b_val): + a_val * (0.0 if tlm_b_val is None else tlm_b_val)) +@pytest.mark.parametrize("a_val", [2.0, -2.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +def test_pos(a_val, tlm_a_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + x = +a + assert x.block_variable.tlm_value == tlm_a_val + + @pytest.mark.parametrize("a_val", [2.0, -2.0]) @pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) def test_neg(a_val, tlm_a_val): From 42b1d607ed9589209b4e08f9aac8504383f9606d Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 17:41:01 +0100 Subject: [PATCH 20/28] Reverse-over-forward AD: DivBlock --- pyadjoint/adjfloat.py | 10 ++++++++++ tests/pyadjoint/test_reverse_over_forward.py | 19 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index d847609b..90e2984c 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -610,6 +610,16 @@ def evaluate_tlm(self, markings=False): )) )) + def solve_tlm(self): + x, = self.get_outputs() + a, b = self.get_dependencies() + terms = [] + if a.tlm_value is not None: + terms.append(a.tlm_value / b.output) + if b.tlm_value is not None: + terms.append((-a.output / (b.output ** 2)) * b.tlm_value) + x.tlm_value = sum_tlm_terms(terms) + def evaluate_hessian(self, markings=False): output = self.get_outputs()[0] hessian_input = output.hessian_value diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index 15ec7c25..b6bfb3b6 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -127,6 +127,25 @@ def test_mul(a_val, tlm_a_val, b_val, tlm_b_val): + a_val * (0.0 if tlm_b_val is None else tlm_b_val)) +@pytest.mark.parametrize("a_val", [2.0, -2.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("b_val", [4.25, -4.25]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125]) +def test_div(a_val, tlm_a_val, b_val, tlm_b_val): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = AdjFloat(b_val) + b.block_variable.tlm_value = AdjFloat(tlm_b_val) + x = (a ** 2) / b + _ = compute_gradient(x.block_variable.tlm_value, (Control(a), Control(b))) + assert np.allclose( + a.block_variable.adj_value, + (2 / b_val) * tlm_a_val - 2 * a_val / (b_val ** 2) * tlm_b_val) + assert np.allclose( + b.block_variable.adj_value, + - 2 * a_val / (b_val ** 2) * tlm_a_val + 2 * (a_val ** 2) / (b_val ** 3) * tlm_b_val) + + @pytest.mark.parametrize("a_val", [2.0, -2.0]) @pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) def test_pos(a_val, tlm_a_val): From b535ae0301a631440a4863a413ff4dda1a3a853b Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 17:50:17 +0100 Subject: [PATCH 21/28] Reverse-over-forward AD: MinBlock and MaxBlock --- pyadjoint/adjfloat.py | 16 ++++++++ tests/pyadjoint/test_reverse_over_forward.py | 42 ++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 90e2984c..66c5fd95 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -305,6 +305,14 @@ def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepar idx = 0 if inputs[0] <= inputs[1] else 1 return tlm_inputs[idx] + def solve_tlm(self): + x, = self.get_outputs() + a, b = self.get_dependencies() + if a.output <= b.output: + x.tlm_value = +a.tlm_value + else: + x.tlm_value = +b.tlm_value + def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): return self.evaluate_adj_component(inputs, hessian_inputs, block_variable, idx, prepared) @@ -327,6 +335,14 @@ def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx, prepar else: return 0. + def solve_tlm(self): + x, = self.get_outputs() + a, b = self.get_dependencies() + if a.output >= b.output: + x.tlm_value = +a.tlm_value + else: + x.tlm_value = +b.tlm_value + def evaluate_tlm_component(self, inputs, tlm_inputs, block_variable, idx, prepared=None): idx = 0 if inputs[0] >= inputs[1] else 1 return tlm_inputs[idx] diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index b6bfb3b6..c27fe9e6 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -39,6 +39,48 @@ def test_log(a_val, tlm_a_val): assert adj_value == -tlm_a_val / (a_val ** 2) +@pytest.mark.parametrize("a_val", [2.0, 3.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("c", [0, 1]) +def test_min_left(a_val, tlm_a_val, c): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = AdjFloat(a_val + c) + x = min(a, b) + assert x.block_variable.tlm_value == tlm_a_val + + +@pytest.mark.parametrize("b_val", [2.0, 3.0]) +@pytest.mark.parametrize("tlm_b_val", [3.5, -3.5]) +def test_min_right(b_val, tlm_b_val): + a = AdjFloat(b_val + 1) + b = AdjFloat(b_val) + b.block_variable.tlm_value = AdjFloat(tlm_b_val) + x = min(a, b) + assert x.block_variable.tlm_value == tlm_b_val + + +@pytest.mark.parametrize("a_val", [2.0, 3.0]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("c", [0, -1]) +def test_max_left(a_val, tlm_a_val, c): + a = AdjFloat(a_val) + a.block_variable.tlm_value = AdjFloat(tlm_a_val) + b = AdjFloat(a_val + c) + x = max(a, b) + assert x.block_variable.tlm_value == tlm_a_val + + +@pytest.mark.parametrize("b_val", [2.0, 3.0]) +@pytest.mark.parametrize("tlm_b_val", [3.5, -3.5]) +def test_max_right(b_val, tlm_b_val): + a = AdjFloat(b_val - 1) + b = AdjFloat(b_val) + b.block_variable.tlm_value = AdjFloat(tlm_b_val) + x = max(a, b) + assert x.block_variable.tlm_value == tlm_b_val + + @pytest.mark.parametrize("a_val", [2.0, 3.0]) @pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) @pytest.mark.parametrize("b_val", [4.25, 5.25]) From 7ff113d5026850270591d9a607e2be36308cb1c1 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Thu, 11 Jul 2024 18:21:15 +0100 Subject: [PATCH 22/28] == -> np.allclose --- tests/pyadjoint/test_reverse_over_forward.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index c27fe9e6..555a787d 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -25,7 +25,7 @@ def test_exp(a_val, tlm_a_val): x = exp(a) _ = compute_gradient(x.block_variable.tlm_value, Control(a)) adj_value = a.block_variable.adj_value - assert adj_value == exp(a_val) * tlm_a_val + assert np.allclose(adj_value, exp(a_val) * tlm_a_val) @pytest.mark.parametrize("a_val", [2.0, 3.0]) @@ -36,7 +36,7 @@ def test_log(a_val, tlm_a_val): x = log(a) _ = compute_gradient(x.block_variable.tlm_value, Control(a)) adj_value = a.block_variable.adj_value - assert adj_value == -tlm_a_val / (a_val ** 2) + assert np.allclose(adj_value, -tlm_a_val / (a_val ** 2)) @pytest.mark.parametrize("a_val", [2.0, 3.0]) From 7deeb81fd8261486525360a6de0b5f490f2d2641 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Fri, 12 Jul 2024 11:18:47 +0100 Subject: [PATCH 23/28] More reverse-over-forward testing --- tests/pyadjoint/test_reverse_over_forward.py | 72 ++++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index 555a787d..a298c09e 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -23,6 +23,7 @@ def test_exp(a_val, tlm_a_val): a = AdjFloat(a_val) a.block_variable.tlm_value = AdjFloat(tlm_a_val) x = exp(a) + stop_annotating() _ = compute_gradient(x.block_variable.tlm_value, Control(a)) adj_value = a.block_variable.adj_value assert np.allclose(adj_value, exp(a_val) * tlm_a_val) @@ -34,6 +35,7 @@ def test_log(a_val, tlm_a_val): a = AdjFloat(a_val) a.block_variable.tlm_value = AdjFloat(tlm_a_val) x = log(a) + stop_annotating() _ = compute_gradient(x.block_variable.tlm_value, Control(a)) adj_value = a.block_variable.adj_value assert np.allclose(adj_value, -tlm_a_val / (a_val ** 2)) @@ -48,6 +50,12 @@ def test_min_left(a_val, tlm_a_val, c): b = AdjFloat(a_val + c) x = min(a, b) assert x.block_variable.tlm_value == tlm_a_val + y = x ** 3 + stop_annotating() + _ = compute_gradient(y.block_variable.tlm_value, Control(a)) + adj_value = a.block_variable.adj_value + assert np.allclose(adj_value, 6 * a_val * tlm_a_val) + @pytest.mark.parametrize("b_val", [2.0, 3.0]) @@ -58,6 +66,11 @@ def test_min_right(b_val, tlm_b_val): b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = min(a, b) assert x.block_variable.tlm_value == tlm_b_val + y = x ** 3 + stop_annotating() + _ = compute_gradient(y.block_variable.tlm_value, Control(b)) + adj_value = b.block_variable.adj_value + assert np.allclose(adj_value, 6 * b_val * tlm_b_val) @pytest.mark.parametrize("a_val", [2.0, 3.0]) @@ -69,6 +82,11 @@ def test_max_left(a_val, tlm_a_val, c): b = AdjFloat(a_val + c) x = max(a, b) assert x.block_variable.tlm_value == tlm_a_val + y = x ** 3 + stop_annotating() + _ = compute_gradient(y.block_variable.tlm_value, Control(a)) + adj_value = a.block_variable.adj_value + assert np.allclose(adj_value, 6 * a_val * tlm_a_val) @pytest.mark.parametrize("b_val", [2.0, 3.0]) @@ -79,6 +97,11 @@ def test_max_right(b_val, tlm_b_val): b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = max(a, b) assert x.block_variable.tlm_value == tlm_b_val + y = x ** 3 + stop_annotating() + _ = compute_gradient(y.block_variable.tlm_value, Control(b)) + adj_value = b.block_variable.adj_value + assert np.allclose(adj_value, 6 * b_val * tlm_b_val) @pytest.mark.parametrize("a_val", [2.0, 3.0]) @@ -103,6 +126,7 @@ def test_a_pow_a(a_val, tlm_a_val): a = AdjFloat(a_val) a.block_variable.tlm_value = AdjFloat(tlm_a_val) x = a ** a + stop_annotating() _ = compute_gradient(x.block_variable.tlm_value, Control(a)) adj_value = a.block_variable.adj_value assert np.allclose( @@ -127,6 +151,18 @@ def test_add(a_val, tlm_a_val, b_val, tlm_b_val): assert (x.block_variable.tlm_value == (0.0 if tlm_a_val is None else tlm_a_val) + (0.0 if tlm_b_val is None else tlm_b_val)) + y = x ** 3 + stop_annotating() + if tlm_a_val is not None or tlm_b_val is not None: + _ = compute_gradient(y.block_variable.tlm_value, (Control(a), Control(b))) + assert np.allclose( + a.block_variable.adj_value, + (6 * a_val + 6 * b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + + (6 * a_val + 6 * b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) + assert np.allclose( + b.block_variable.adj_value, + (6 * a_val + 6 * b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + + (6 * a_val + 6 * b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) @@ -147,6 +183,18 @@ def test_sub(a_val, tlm_a_val, b_val, tlm_b_val): assert (x.block_variable.tlm_value == (0.0 if tlm_a_val is None else tlm_a_val) - (0.0 if tlm_b_val is None else tlm_b_val)) + y = x ** 3 + stop_annotating() + if tlm_a_val is not None or tlm_b_val is not None: + _ = compute_gradient(y.block_variable.tlm_value, (Control(a), Control(b))) + assert np.allclose( + a.block_variable.adj_value, + (6 * a_val - 6 * b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + + (-6 * a_val + 6 * b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) + assert np.allclose( + b.block_variable.adj_value, + (-6 * a_val + 6 * b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + + (6 * a_val - 6 * b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) @@ -167,6 +215,19 @@ def test_mul(a_val, tlm_a_val, b_val, tlm_b_val): assert (x.block_variable.tlm_value == b_val * (0.0 if tlm_a_val is None else tlm_a_val) + a_val * (0.0 if tlm_b_val is None else tlm_b_val)) + stop_annotating() + if tlm_a_val is not None or tlm_b_val is not None: + _ = compute_gradient(x.block_variable.tlm_value, (Control(a), Control(b))) + if tlm_b_val is None: + assert a.block_variable.adj_value is None + else: + assert np.allclose( + a.block_variable.adj_value, tlm_b_val) + if tlm_a_val is None: + assert b.block_variable.adj_value is None + else: + assert np.allclose( + b.block_variable.adj_value, tlm_a_val) @pytest.mark.parametrize("a_val", [2.0, -2.0]) @@ -179,6 +240,7 @@ def test_div(a_val, tlm_a_val, b_val, tlm_b_val): b = AdjFloat(b_val) b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = (a ** 2) / b + stop_annotating() _ = compute_gradient(x.block_variable.tlm_value, (Control(a), Control(b))) assert np.allclose( a.block_variable.adj_value, @@ -195,6 +257,11 @@ def test_pos(a_val, tlm_a_val): a.block_variable.tlm_value = AdjFloat(tlm_a_val) x = +a assert x.block_variable.tlm_value == tlm_a_val + y = x ** 3 + stop_annotating() + _ = compute_gradient(y.block_variable.tlm_value, Control(a)) + adj_value = a.block_variable.adj_value + assert np.allclose(adj_value, 6 * a_val * tlm_a_val) @pytest.mark.parametrize("a_val", [2.0, -2.0]) @@ -204,3 +271,8 @@ def test_neg(a_val, tlm_a_val): a.block_variable.tlm_value = AdjFloat(tlm_a_val) x = -a assert x.block_variable.tlm_value == -tlm_a_val + y = x ** 3 + stop_annotating() + _ = compute_gradient(y.block_variable.tlm_value, Control(a)) + adj_value = a.block_variable.adj_value + assert np.allclose(adj_value, -6 * a_val * tlm_a_val) From 8a9c334dd8738e67f1c9d71db87b10be15de4f86 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Fri, 12 Jul 2024 11:56:53 +0100 Subject: [PATCH 24/28] Bugfix --- pyadjoint/adjfloat.py | 2 +- tests/pyadjoint/test_reverse_over_forward.py | 39 +++++++++----------- 2 files changed, 18 insertions(+), 23 deletions(-) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index 66c5fd95..ac4536e4 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -431,7 +431,7 @@ def solve_tlm(self): if a.tlm_value is not None: terms.append(b.output * (a.output ** (b.output - 1)) * a.tlm_value) if b.tlm_value is not None: - terms.append(log(a.output) * (a.output ** b.output) * a.tlm_value) + terms.append(log(a.output) * (a.output ** b.output) * b.tlm_value) x.tlm_value = sum_tlm_terms(terms) def evaluate_hessian(self, markings=False): diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index a298c09e..74a889fb 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -105,33 +105,28 @@ def test_max_right(b_val, tlm_b_val): @pytest.mark.parametrize("a_val", [2.0, 3.0]) -@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5, None]) @pytest.mark.parametrize("b_val", [4.25, 5.25]) -@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125, None]) def test_pow(a_val, tlm_a_val, b_val, tlm_b_val): a = AdjFloat(a_val) - a.block_variable.tlm_value = AdjFloat(tlm_a_val) + if tlm_a_val is not None: + a.block_variable.tlm_value = AdjFloat(tlm_a_val) b = AdjFloat(b_val) - b.block_variable.tlm_value = AdjFloat(tlm_b_val) + if tlm_b_val is not None: + b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = a ** b - J_hat = ReducedFunctional(x.block_variable.tlm_value, Control(a)) - assert taylor_test(J_hat, a, AdjFloat(1.0)) > 1.9 - J_hat = ReducedFunctional(x.block_variable.tlm_value, Control(b)) - assert taylor_test(J_hat, b, AdjFloat(1.0)) > 1.9 - - -@pytest.mark.parametrize("a_val", [2.0, 3.0]) -@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) -def test_a_pow_a(a_val, tlm_a_val): - a = AdjFloat(a_val) - a.block_variable.tlm_value = AdjFloat(tlm_a_val) - x = a ** a - stop_annotating() - _ = compute_gradient(x.block_variable.tlm_value, Control(a)) - adj_value = a.block_variable.adj_value - assert np.allclose( - adj_value, ((1 + log(a_val)) ** 2 + (1 / a_val)) * (a_val ** a_val) * tlm_a_val) - + if tlm_a_val is not None or tlm_b_val is not None: + _ = compute_gradient(x.block_variable.tlm_value, (Control(a), Control(b))) + assert np.allclose( + a.block_variable.adj_value, + b_val * (b_val - 1) * (a_val ** (b_val - 2)) * (0.0 if tlm_a_val is None else tlm_a_val) + + (1 + b_val * log(a_val)) * (a_val ** (b_val - 1)) * (0.0 if tlm_b_val is None else tlm_b_val)) + assert np.allclose( + b.block_variable.adj_value, + (1 + b_val * log(a_val)) * (a_val ** (b_val - 1)) * (0.0 if tlm_a_val is None else tlm_a_val) + + (log(a_val) ** 2) * (a_val ** b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) + @pytest.mark.parametrize("a_val", [2.0, -2.0]) @pytest.mark.parametrize("tlm_a_val", [3.5, -3.5, None]) From 1d84ccf1c63374d8c3096ce82e66624f6c634b75 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Fri, 12 Jul 2024 12:02:17 +0100 Subject: [PATCH 25/28] Extra AdjFloat.__truediv__ testing --- tests/pyadjoint/test_reverse_over_forward.py | 27 ++++++++++++-------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index 74a889fb..03544cab 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -226,23 +226,28 @@ def test_mul(a_val, tlm_a_val, b_val, tlm_b_val): @pytest.mark.parametrize("a_val", [2.0, -2.0]) -@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5]) +@pytest.mark.parametrize("tlm_a_val", [3.5, -3.5, None]) @pytest.mark.parametrize("b_val", [4.25, -4.25]) -@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125]) +@pytest.mark.parametrize("tlm_b_val", [5.8125, -5.8125, None]) def test_div(a_val, tlm_a_val, b_val, tlm_b_val): a = AdjFloat(a_val) - a.block_variable.tlm_value = AdjFloat(tlm_a_val) + if tlm_a_val is not None: + a.block_variable.tlm_value = AdjFloat(tlm_a_val) b = AdjFloat(b_val) - b.block_variable.tlm_value = AdjFloat(tlm_b_val) + if tlm_b_val is not None: + b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = (a ** 2) / b stop_annotating() - _ = compute_gradient(x.block_variable.tlm_value, (Control(a), Control(b))) - assert np.allclose( - a.block_variable.adj_value, - (2 / b_val) * tlm_a_val - 2 * a_val / (b_val ** 2) * tlm_b_val) - assert np.allclose( - b.block_variable.adj_value, - - 2 * a_val / (b_val ** 2) * tlm_a_val + 2 * (a_val ** 2) / (b_val ** 3) * tlm_b_val) + if tlm_a_val is not None or tlm_b_val is not None: + _ = compute_gradient(x.block_variable.tlm_value, (Control(a), Control(b))) + assert np.allclose( + a.block_variable.adj_value, + (2 / b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + - 2 * a_val / (b_val ** 2) * (0.0 if tlm_b_val is None else tlm_b_val)) + assert np.allclose( + b.block_variable.adj_value, + - 2 * a_val / (b_val ** 2) * (0.0 if tlm_a_val is None else tlm_a_val) + + 2 * (a_val ** 2) / (b_val ** 3) * (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) From 8fa8a0abf0b3f621a343e72b9a2feefdc22cd207 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Fri, 12 Jul 2024 12:14:01 +0100 Subject: [PATCH 26/28] Minor __pos__ fixes --- pyadjoint/adjfloat.py | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/pyadjoint/adjfloat.py b/pyadjoint/adjfloat.py index ac4536e4..8f303348 100644 --- a/pyadjoint/adjfloat.py +++ b/pyadjoint/adjfloat.py @@ -378,15 +378,6 @@ def __str__(self): return f"{self.terms[0]} {self.symbol} {self.terms[1]}" -def sum_tlm_terms(terms): - if len(terms) == 0: - return None - elif len(terms) == 1: - return +terms[0] - else: - return sum(terms[1:], start=terms[0]) - - class PowBlock(FloatOperatorBlock): operator = staticmethod(float.__pow__) symbol = "**" @@ -432,7 +423,7 @@ def solve_tlm(self): terms.append(b.output * (a.output ** (b.output - 1)) * a.tlm_value) if b.tlm_value is not None: terms.append(log(a.output) * (a.output ** b.output) * b.tlm_value) - x.tlm_value = sum_tlm_terms(terms) + x.tlm_value = None if len(terms) == 0 else sum(terms[1:], start=terms[0]) def evaluate_hessian(self, markings=False): output = self.get_outputs()[0] @@ -501,7 +492,12 @@ def solve_tlm(self): x, = self.get_outputs() terms = tuple(dep.tlm_value for dep in self.get_dependencies() if dep.tlm_value is not None) - x.tlm_value = sum_tlm_terms(terms) + if len(terms) == 0: + x.tlm_value = None + elif len(terms) == 1: + x.tlm_value = +terms[0] + else: + x.tlm_value = sum(terms[1:], start=terms[0]) def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): @@ -536,7 +532,7 @@ def solve_tlm(self): else: x.tlm_value = -b.tlm_value elif b.tlm_value is None: - x.tlm_value = a.tlm_value + x.tlm_value = +a.tlm_value else: x.tlm_value = a.tlm_value - b.tlm_value @@ -576,7 +572,7 @@ def solve_tlm(self): terms.append(b.output * a.tlm_value) if b.tlm_value is not None: terms.append(a.output * b.tlm_value) - x.tlm_value = sum_tlm_terms(terms) + x.tlm_value = None if len(terms) == 0 else sum(terms[1:], start=terms[0]) def evaluate_hessian_component(self, inputs, hessian_inputs, adj_inputs, block_variable, idx, relevant_dependencies, prepared=None): @@ -634,7 +630,7 @@ def solve_tlm(self): terms.append(a.tlm_value / b.output) if b.tlm_value is not None: terms.append((-a.output / (b.output ** 2)) * b.tlm_value) - x.tlm_value = sum_tlm_terms(terms) + x.tlm_value = None if len(terms) == 0 else sum(terms[1:], start=terms[0]) def evaluate_hessian(self, markings=False): output = self.get_outputs()[0] From 05a8365370097790739ceb624b3506f26bd00231 Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Fri, 12 Jul 2024 12:24:56 +0100 Subject: [PATCH 27/28] Documentation fixes --- pyadjoint/block_variable.py | 2 +- pyadjoint/tape.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/pyadjoint/block_variable.py b/pyadjoint/block_variable.py index 9c4ae10c..bffd9c2e 100644 --- a/pyadjoint/block_variable.py +++ b/pyadjoint/block_variable.py @@ -102,7 +102,7 @@ def restore_output(self): the value of `self.output` to `self.block_variable.saved_output`. Returns: - The context manager + The context manager. """ if self.output is self.saved_output: diff --git a/pyadjoint/tape.py b/pyadjoint/tape.py index b1668ded..f7944dcb 100644 --- a/pyadjoint/tape.py +++ b/pyadjoint/tape.py @@ -160,11 +160,11 @@ def continue_reverse_over_forward(): @contextmanager def stop_reverse_over_forward(): - """Return a callable used to construct a context manager within which - reverse-over-forward AD is disabled. + """Return a context manager used to temporarily disable + reverse-over-forward AD. Returns: - callable: Callable which returns a context manager. + The context manager. """ global _reverse_over_forward_enabled @@ -177,7 +177,8 @@ def stop_reverse_over_forward(): def no_reverse_over_forward(function): - """Decorator to disable reverse-over-forward AD for the decorated callable. + """Decorator to temporarily disable reverse-over-forward AD for the + decorated callable. Args: function (callable): The callable. From ec6e2d5793a21a5a499a4102c98989b7390f6c1a Mon Sep 17 00:00:00 2001 From: "James R. Maddison" Date: Fri, 12 Jul 2024 12:42:28 +0100 Subject: [PATCH 28/28] Test updates --- tests/pyadjoint/test_reverse_over_forward.py | 28 ++++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/tests/pyadjoint/test_reverse_over_forward.py b/tests/pyadjoint/test_reverse_over_forward.py index 03544cab..9d086a52 100644 --- a/tests/pyadjoint/test_reverse_over_forward.py +++ b/tests/pyadjoint/test_reverse_over_forward.py @@ -116,6 +116,12 @@ def test_pow(a_val, tlm_a_val, b_val, tlm_b_val): if tlm_b_val is not None: b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = a ** b + if tlm_a_val is None and tlm_b_val is None: + assert x.block_variable.tlm_value is None + else: + assert (x.block_variable.tlm_value == + b_val * (a_val ** (b_val - 1)) * (0.0 if tlm_a_val is None else tlm_a_val) + + log(a_val) * (a_val ** b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) if tlm_a_val is not None or tlm_b_val is not None: _ = compute_gradient(x.block_variable.tlm_value, (Control(a), Control(b))) assert np.allclose( @@ -152,12 +158,12 @@ def test_add(a_val, tlm_a_val, b_val, tlm_b_val): _ = compute_gradient(y.block_variable.tlm_value, (Control(a), Control(b))) assert np.allclose( a.block_variable.adj_value, - (6 * a_val + 6 * b_val) * (0.0 if tlm_a_val is None else tlm_a_val) - + (6 * a_val + 6 * b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) + 6 * (a_val + b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + + 6 * (a_val + b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) assert np.allclose( b.block_variable.adj_value, - (6 * a_val + 6 * b_val) * (0.0 if tlm_a_val is None else tlm_a_val) - + (6 * a_val + 6 * b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) + 6 * (a_val + b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + + 6 * (a_val + b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) @@ -184,12 +190,12 @@ def test_sub(a_val, tlm_a_val, b_val, tlm_b_val): _ = compute_gradient(y.block_variable.tlm_value, (Control(a), Control(b))) assert np.allclose( a.block_variable.adj_value, - (6 * a_val - 6 * b_val) * (0.0 if tlm_a_val is None else tlm_a_val) - + (-6 * a_val + 6 * b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) + 6 * (a_val - b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + - 6 * (a_val - b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) assert np.allclose( b.block_variable.adj_value, - (-6 * a_val + 6 * b_val) * (0.0 if tlm_a_val is None else tlm_a_val) - + (6 * a_val - 6 * b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) + - 6 * (a_val - b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + + 6 * (a_val - b_val) * (0.0 if tlm_b_val is None else tlm_b_val)) @pytest.mark.parametrize("a_val", [2.0, -2.0]) @@ -237,6 +243,12 @@ def test_div(a_val, tlm_a_val, b_val, tlm_b_val): if tlm_b_val is not None: b.block_variable.tlm_value = AdjFloat(tlm_b_val) x = (a ** 2) / b + if tlm_a_val is None and tlm_b_val is None: + assert x.block_variable.tlm_value is None + else: + assert (x.block_variable.tlm_value == + (2 * a_val / b_val) * (0.0 if tlm_a_val is None else tlm_a_val) + - ((a_val ** 2) / (b_val ** 2)) * (0.0 if tlm_b_val is None else tlm_b_val)) stop_annotating() if tlm_a_val is not None or tlm_b_val is not None: _ = compute_gradient(x.block_variable.tlm_value, (Control(a), Control(b)))