Optimize optimize() traversal with memoization and fixed-point rewrite loops

CHANGELOG.md

@@ -7,6 +7,9 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 ## Version 0.4.6 - Unreleased
 
+### Performance
+* Improve optimize() performance via per-call memoization, reduced allocations, and fixed-point rewrite loops; no behavior change intended.
+
 ### Fix
 * Handle case when input sensorchan strings are string subclasses.
 * Fix issue where lazy warps did not respect explicitly given dsize arguments

delayed_image/delayed_base.py

@@ -1,6 +1,7 @@
 """
 Abstract nodes
 """
+from __future__ import annotations
 import numpy as np
 import ubelt as ub
 
@@ -13,6 +14,18 @@
 USE_SLOTS = True
 
 
+# Per-call optimization context
+class OptimizeContext:
+    """
+    Holds per-call optimization state to avoid repeated work.
+    """
+    if USE_SLOTS:
+        __slots__ = ('memo',)
+
+    def __init__(self):
+        self.memo = {}
+
+
 # from kwcoco.util.util_monkey import Reloadable  # NOQA
 # @Reloadable.developing  # NOQA
 class DelayedOperation:
@@ -385,7 +398,7 @@ def finalize(self, prepare=True, optimize=True, **kwargs):
         # final = np.asanyarray(final) # does not work with xarray
         return final
 
-    def optimize(self):
+    def optimize(self, ctx=None):
         """
         Returns:
             DelayedOperation

delayed_image/delayed_base.pyi

@@ -9,6 +9,13 @@ from _typeshed import Incomplete
 from collections.abc import Generator
 
 
+class OptimizeContext:
+    memo: Dict[int, 'DelayedOperation']
+
+    def __init__(self) -> None:
+        ...
+
+
 class DelayedOperation(ub.NiceRepr):
     meta: Incomplete
 
@@ -57,7 +64,7 @@ class DelayedOperation(ub.NiceRepr):
                  **kwargs) -> ArrayLike:
         ...
 
-    def optimize(self) -> DelayedOperation:
+    def optimize(self, ctx: OptimizeContext | None = None) -> DelayedOperation:
         ...
 
 

delayed_image/delayed_leafs.py

@@ -1,6 +1,7 @@
 """
 Terminal nodes
 """
+from __future__ import annotations
 
 import kwarray
 import kwimage
@@ -30,9 +31,15 @@ def get_transform_from_leaf(self):
         """
         return kwimage.Affine.eye()
 
-    def optimize(self):
+    def optimize(self, ctx=None):
+        if ctx is None:
+            ctx = delayed_base.OptimizeContext()
+        memo = ctx.memo
+        if self in memo:
+            return memo[self]
         if TRACE_OPTIMIZE:
             self._opt_logs.append('optimize DelayedImageLeaf')
+        memo[self] = self
         return self
 
 

delayed_image/delayed_leafs.pyi

@@ -3,6 +3,7 @@ from os import PathLike
 from typing import Tuple
 from _typeshed import Incomplete
 from delayed_image.delayed_nodes import DelayedImage
+from delayed_image.delayed_base import OptimizeContext
 
 from delayed_image.channel_spec import FusedChannelSpec
 
@@ -14,7 +15,7 @@ class DelayedImageLeaf(DelayedImage):
     def get_transform_from_leaf(self) -> kwimage.Affine:
         ...
 
-    def optimize(self):
+    def optimize(self, ctx: OptimizeContext | None = None):
         ...