Fix: preserve direction and input identity in memoized evaluations

examples/plot_memoizeoperator.py

@@ -0,0 +1,46 @@
+r"""
+Memoized forward and adjoint evaluations
+========================================
+
+This example shows how to use the :class:`pylops.MemoizeOperator`
+operator to cache repeated applications of an operator.
+"""
+
+import numpy as np
+
+import pylops
+
+###############################################################################
+# Let's start by creating a :class:`pylops.Diagonal` operator and wrapping it
+# into a :class:`pylops.MemoizeOperator` operator choosing ``max_neval=2``
+# (meaning that at most 4 evaluations are cached).
+Op = pylops.Diagonal(np.array([2.0, 3.0, 4.0]))
+MOp = pylops.MemoizeOperator(Op, max_neval=4)
+
+###############################################################################
+# Now we first apply twice the forward pass of operator (and check that only
+# one entry is memoized), followed by one adjoint pass of operator (and check
+# that now two entries are memoized)
+
+x = np.array([1.0, -2.0, 0.5])
+y = MOp @ x
+y = MOp @ x
+print("Number of evaluations: ", MOp.neval)
+
+xadj = MOp.H @ y
+print("Number of evaluations: ", MOp.neval)
+
+###############################################################################
+# Repeating either calculation reuses its own cached result.
+print("y - ycached = ", y - MOp @ x)
+print("xadj - xadjcached = ", xadj - MOp.H @ y)
+print("Number of evaluations: ", MOp.neval)
+
+###############################################################################
+# Finally, we can see that similar inputs are not interchangeable: any
+# small change triggers a new memoization.
+
+x_nearby = x.copy()
+x_nearby[0] += 1e-6
+y_nearby = MOp @ x_nearby
+print("Number of evaluations: ", MOp.neval)

pylops/basicoperators/memoizeoperator.py

@@ -11,9 +11,10 @@ class MemoizeOperator(LinearOperator):
     r"""Memoize Operator.
 
     This operator can be used to wrap any PyLops operator and add a memoize
-    functionality
-    and stores the last ``max_neval`` model/data
-    vector pairs
+    functionality - i.e., store the last ``max_neval`` model/data
+    vector pairs for forward and adjoint evaluations. Note that cache
+    hits require identical input values and dtype; nearby vectors are
+    evaluated independently.
 
     Parameters
     ----------
@@ -43,26 +44,39 @@ def __init__(
 
         self.max_neval = max_neval
         self.store: list[tuple[NDArray, NDArray]] = []  # Store a list of (x, y)
+        self._store_is_adjoint: list[bool] = []
         self.neval = 0  # Number of evaluations of the operator
 
     def _matvec(self, x: NDArray) -> NDArray:
-        for xstored, ystored in self.store:
-            if np.allclose(xstored, x):
+        for (xstored, ystored), is_adjoint in zip(
+            self.store, self._store_is_adjoint, strict=True
+        ):
+            if (
+                not is_adjoint
+                and xstored.dtype == x.dtype
+                and np.array_equal(xstored, x)
+            ):
                 return ystored
         if len(self.store) + 1 > self.max_neval:
             del self.store[0]  # Delete oldest
+            del self._store_is_adjoint[0]
         y = self.Op._matvec(x)
         self.neval += 1
         self.store.append((x.copy(), y.copy()))
+        self._store_is_adjoint.append(False)
         return y
 
     def _rmatvec(self, y: NDArray) -> NDArray:
-        for xstored, ystored in self.store:
-            if np.allclose(ystored, y):
+        for (xstored, ystored), is_adjoint in zip(
+            self.store, self._store_is_adjoint, strict=True
+        ):
+            if is_adjoint and ystored.dtype == y.dtype and np.array_equal(ystored, y):
                 return xstored
         if len(self.store) + 1 > self.max_neval:
             del self.store[0]  # Delete oldest
+            del self._store_is_adjoint[0]
         x = self.Op._rmatvec(y)
         self.neval += 1
         self.store.append((x.copy(), y.copy()))
+        self._store_is_adjoint.append(True)
         return x

pytests/test_memoizeoperator.py

@@ -2,25 +2,28 @@
 
 if int(os.environ.get("TEST_CUPY_PYLOPS", 0)):
     import cupy as np
-    from cupy.testing import assert_array_almost_equal
+    from cupy.testing import assert_array_almost_equal, assert_array_equal
 
     backend = "cupy"
 else:
     import numpy as np
-    from numpy.testing import assert_array_almost_equal
+    from numpy.testing import assert_array_almost_equal, assert_array_equal
 
     backend = "numpy"
 import pytest
 
 from pylops import MemoizeOperator
-from pylops.basicoperators import MatrixMult, VStack
+from pylops.basicoperators import Diagonal, MatrixMult, VStack
 
 par1 = {"ny": 11, "nx": 11, "imag": 0, "dtype": "float32"}  # square real
 par1j = {"ny": 11, "nx": 11, "imag": 1j, "dtype": "complex64"}  # square imag
+par2 = {"ny": 11, "nx": 7, "imag": 0, "dtype": "float32"}  # rectangular real
+par2j = {"ny": 11, "nx": 7, "imag": 1j, "dtype": "complex64"}  # rectangular imag
 
 
-@pytest.mark.parametrize("par", [(par1), (par1j)])
-def test_memoize_evals(par):
+@pytest.mark.parametrize("par", [(par1), (par1j), (par2), (par2j)])
+@pytest.mark.parametrize("adjoint", [False, True])
+def test_memoize_evals(par, adjoint):
     """Check nevals counter when same model/data vectors are inputted
     to the operator
     """
@@ -31,28 +34,95 @@ def test_memoize_evals(par):
     Aop = MatrixMult(A, dtype=par["dtype"])
     Amemop = MemoizeOperator(Aop, max_neval=2)
 
+    if adjoint:
+        Aop = Aop.H
+        Amemop1 = Amemop.H
+        n = par["ny"]
+    else:
+        Amemop1 = Amemop
+        n = par["nx"]
+
     # 1st evaluation
-    Amemop * np.ones(par["nx"])
+    yref = Aop * np.ones(n)
+    y = Amemop1 * np.ones(n)
     assert Amemop.neval == 1
     # repeat 1st evaluation multiple times
     for _ in range(2):
-        Amemop * np.ones(par["nx"])
+        y1 = Amemop1 * np.ones(n)
     assert Amemop.neval == 1
+    assert_array_equal(yref, y)
+    assert_array_equal(yref, y1)
     # 2nd evaluation
-    Amemop * np.full(par["nx"], 2)  # same
+    Amemop1 * np.full(n, 2)  # same
     assert Amemop.neval == 2
     # 3rd evaluation (np.ones goes out of store)
-    Amemop * np.full(par["nx"], 3)  # same
+    Amemop1 * np.full(n, 3)  # same
     assert Amemop.neval == 3
     # 4th evaluation
-    Amemop * np.ones(par["nx"])
+    Amemop1 * np.ones(n)
     assert Amemop.neval == 4
 
 
+@pytest.mark.parametrize("par", [(par1), (par2)])
+@pytest.mark.parametrize("adjoint", [False, True])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_memoize_nearby_inputs(par, adjoint, dtype):
+    """Check that nearby vectors (even within allclose)
+    are identified as distinct evaluations."""
+    Op = Diagonal(np.arange(par["nx"], dtype=dtype), dtype=dtype)
+    MOp = MemoizeOperator(Op)
+
+    Op = Op.H if adjoint else Op
+    M1Op = MOp.H if adjoint else MOp
+    x = np.ones(par["nx"], dtype=dtype)
+    x_nearby = x.copy()
+    x_nearby[0] += 1e-6
+    assert np.allclose(x, x_nearby)
+
+    # 2 evals with original and nearby input
+    M1Op @ x
+    y = M1Op @ x_nearby
+
+    assert MOp.neval == 2
+    assert_array_equal(y, Op @ x_nearby)
+    assert_array_equal(y, M1Op @ x_nearby)
+
+
+@pytest.mark.parametrize("par", [(par1)])
+@pytest.mark.parametrize("adjoint", [False, True])
+def test_memoize_input_dtype(par, adjoint):
+    """Check that equal vectors with different precision
+    are identified as distinct evaluations."""
+    low = np.complex64 if par["imag"] == 1j else np.float32
+    high = np.complex128 if par["imag"] == 1j else np.float64
+    A = np.random.normal(0, 10, (par["ny"], par["nx"])).astype("float32") + par[
+        "imag"
+    ] * np.random.normal(0, 10, (par["ny"], par["nx"])).astype("float32")
+
+    Aop = MatrixMult(A, dtype=low)
+    Amemop = MemoizeOperator(Aop)
+
+    Aop = Aop.H if adjoint else Aop
+    Amemop1 = Amemop.H if adjoint else Amemop
+
+    x = np.ones(par["nx"], dtype=low)
+    x_promoted = x.astype(high)
+
+    # 2 evals with original and promoted input
+    Amemop1 @ x
+    result = Amemop1 @ x_promoted
+    expected = Aop @ x_promoted
+
+    assert result.dtype == expected.dtype
+    assert_array_equal(result, expected)
+    assert Amemop.neval == 2
+
+
 @pytest.mark.parametrize(
     "par",
     [
         (par1j),
+        (par2j),
     ],
 )
 def test_memoize_evals_2(par):