From 86197a887d9daf366e895cd9c9e5eee4b48e3bc7 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 11 Oct 2024 10:54:16 +0100 Subject: [PATCH 1/6] ReducedFunctional.optimize_tape will optionally replace self.tape with optimized tape, but not modify original tape --- pyadjoint/reduced_functional.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pyadjoint/reduced_functional.py b/pyadjoint/reduced_functional.py index 8b2af45f..906d670a 100644 --- a/pyadjoint/reduced_functional.py +++ b/pyadjoint/reduced_functional.py @@ -224,7 +224,9 @@ def __call__(self, values): return func_value - def optimize_tape(self): + def optimize_tape(self, replace=False): + if replace: + self.tape = self.tape.copy() self.tape.optimize( controls=self.controls, functionals=[self.functional] From 58466321e3bf7a282f2783fcfd9ddce195641500 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 11 Oct 2024 11:55:58 +0100 Subject: [PATCH 2/6] ReducedFunctional.optimize_tape returns optimized tape --- pyadjoint/reduced_functional.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyadjoint/reduced_functional.py b/pyadjoint/reduced_functional.py index 906d670a..6dccd88b 100644 --- a/pyadjoint/reduced_functional.py +++ b/pyadjoint/reduced_functional.py @@ -231,6 +231,7 @@ def optimize_tape(self, replace=False): controls=self.controls, functionals=[self.functional] ) + return self.tape def marked_controls(self): return marked_controls(self) From 71ed944d631a3be693797745b378cb86d0f13fb5 Mon Sep 17 00:00:00 2001 From: JHopeCollins Date: Thu, 24 Oct 2024 10:28:31 +0100 Subject: [PATCH 3/6] ReducedFunctional uses an optimized copy of the tape --- pyadjoint/reduced_functional.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/pyadjoint/reduced_functional.py b/pyadjoint/reduced_functional.py index 6dccd88b..9151ff98 100644 --- a/pyadjoint/reduced_functional.py +++ b/pyadjoint/reduced_functional.py @@ -77,7 +77,6 @@ def __init__(self, functional, controls, if not isinstance(functional, OverloadedType): raise TypeError("Functional must be an OverloadedType.") self.functional = functional - self.tape = get_working_tape() if tape is None else tape self.controls = Enlist(controls) self.derivative_components = derivative_components self.scale = scale @@ -88,6 +87,13 @@ def __init__(self, functional, controls, self.hessian_cb_pre = hessian_cb_pre self.hessian_cb_post = hessian_cb_post + tape = get_working_tape() if tape is None else tape + self.tape = tape.copy() + self.tape.optimize( + controls=self.controls, + functionals=[self.functional] + ) + if self.derivative_components: # pre callback self.derivative_cb_pre = _get_extract_derivative_components( @@ -224,13 +230,9 @@ def __call__(self, values): return func_value - def optimize_tape(self, replace=False): - if replace: - self.tape = self.tape.copy() - self.tape.optimize( - controls=self.controls, - functionals=[self.functional] - ) + def optimize_tape(self): + # Tape already optimized in __init__ + # TODO: What should we do here now? return self.tape def marked_controls(self): From 3ba7929d61f5a59a5e001bff866dfec3c5fedaa5 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 24 Oct 2024 20:23:53 +0100 Subject: [PATCH 4/6] should copying a control give the current or original value? --- pyadjoint/control.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyadjoint/control.py b/pyadjoint/control.py index f8e626be..ff4f18c6 100644 --- a/pyadjoint/control.py +++ b/pyadjoint/control.py @@ -80,6 +80,7 @@ def assign_numpy(self, dst, src, offset): def fetch_numpy(self, value): return self.control._ad_to_list(value) + # TODO: This should be self.block_variable.checkpoint? def copy_data(self): return self.control._ad_copy() From 1f728f54ad60e4ff12c1ec28f660ed97f6fc0771 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Thu, 24 Oct 2024 20:25:19 +0100 Subject: [PATCH 5/6] Tape.copy now also copies over the checkpoint manager and progress bar, and just transfers any package data that doesn't have a copy method --- pyadjoint/tape.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/pyadjoint/tape.py b/pyadjoint/tape.py index 06a49782..cade09a9 100644 --- a/pyadjoint/tape.py +++ b/pyadjoint/tape.py @@ -368,10 +368,24 @@ def copy(self): """ # TODO: Offer deepcopying. But is it feasible memory wise to copy all checkpoints? - return Tape( + + # TODO: firedrake.DiskCheckPointer doesn't implement + # copy so we have to do something else here. + package_data = {} + for k, v in self._package_data.items(): + try: + package_data[k] = v.copy() + except NotImplementedError: + package_data[k] = v + tape = Tape( blocks=self._blocks, - package_data={k: v.copy() for k, v in self._package_data.items()} + package_data=package_data ) + if self._checkpoint_manager is not None: + tape._checkpoint_manager = self._checkpoint_manager + if self._bar is not _NullProgressBar: + tape.progress_bar = self.progress_bar + return tape def checkpoint_block_vars(self, controls=[], tag=None): """Returns an object to checkpoint the current state of all block variables on the tape. From 87d1ea4fe83c28ae6f140e7a350dcf0b6fbf8484 Mon Sep 17 00:00:00 2001 From: Josh Hope-Collins Date: Fri, 25 Oct 2024 13:48:54 +0100 Subject: [PATCH 6/6] tape.copy copies over latest checkpoint --- pyadjoint/tape.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pyadjoint/tape.py b/pyadjoint/tape.py index cade09a9..dde8d467 100644 --- a/pyadjoint/tape.py +++ b/pyadjoint/tape.py @@ -383,6 +383,7 @@ def copy(self): ) if self._checkpoint_manager is not None: tape._checkpoint_manager = self._checkpoint_manager + tape.latest_checkpoint = self.latest_checkpoint if self._bar is not _NullProgressBar: tape.progress_bar = self.progress_bar return tape