Fix: preserve direction and input identity in memoized evaluations

examples/plot_memoizeoperator.py

@@ -0,0 +1,39 @@
+r"""
+Memoized forward and adjoint evaluations
+=======================================
+
+:class:`pylops.MemoizeOperator` caches repeated applications of an operator.
+The forward and adjoint directions have separate cache entries, since an
+adjoint generally does not undo a forward application.
+"""
+
+import numpy as np
+
+import pylops
+
+###############################################################################
+# A non-unitary diagonal operator makes the distinction explicit:
+# applying the adjoint to ``A @ x`` gives ``A.H @ A @ x``, rather than ``x``.
+op = pylops.Diagonal(np.array([2.0, 3.0, 4.0]))
+memo = pylops.MemoizeOperator(op, max_neval=2)
+x = np.array([1.0, -2.0, 0.5])
+y = memo @ x
+z = memo.H @ y
+np.testing.assert_allclose(z, op.H @ (op @ x))
+assert memo.neval == 2
+
+###############################################################################
+# Repeating either calculation reuses its own cached result. Both entries
+# share the two-evaluation budget, and the stored vectors remain (x, y) pairs.
+np.testing.assert_allclose(memo @ x, y)
+np.testing.assert_allclose(memo.H @ y, z)
+assert memo.neval == 2
+assert len(memo.store) == 2
+
+
+###############################################################################
+# Similar inputs are not interchangeable: small changes must reach the operator.
+x_nearby = x.copy()
+x_nearby[0] += 1e-6
+np.testing.assert_array_equal(memo @ x_nearby, op @ x_nearby)
+assert memo.neval == 3

pylops/basicoperators/memoizeoperator.py

@@ -13,7 +13,10 @@ class MemoizeOperator(LinearOperator):
     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
+    vector pairs. Forward and adjoint evaluations are cached separately,
+    sharing the same ``max_neval`` limit. An adjoint is not assumed to be
+    the inverse of a forward evaluation. Cache hits require identical input
+    values and dtype; nearby vectors are evaluated independently.
 
     Parameters
     ----------
@@ -43,26 +46,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

@@ -13,7 +13,8 @@
 import pytest
 
 from pylops import MemoizeOperator
-from pylops.basicoperators import MatrixMult, VStack
+from pylops.basicoperators import Diagonal, MatrixMult, VStack
+from pylops.optimization.basic import cgls
 
 par1 = {"ny": 11, "nx": 11, "imag": 0, "dtype": "float32"}  # square real
 par1j = {"ny": 11, "nx": 11, "imag": 1j, "dtype": "complex64"}  # square imag
@@ -80,3 +81,115 @@ def test_memoize_evals_2(par):
     y2 = np.concatenate([np.real(y), np.imag(y)])
     xinv2 = Aop2.div(y2)
     assert_array_almost_equal(xinv2, x, decimal=2 if rdtype == np.float32 else 4)
+
+
+@pytest.mark.parametrize("shape", [(2, 2), (3, 2), (2, 3)])
+@pytest.mark.parametrize("complex_values", [False, True])
+@pytest.mark.parametrize("adjoint_first", [False, True])
+def test_memoize_direction(shape, complex_values, adjoint_first):
+    """Forward/adjoint composition must not reuse an inverse result."""
+    matrix = np.arange(1, shape[0] * shape[1] + 1, dtype=np.float64).reshape(shape)
+    if complex_values:
+        matrix = matrix + 1j * matrix[::-1]
+    op = MatrixMult(matrix, dtype=matrix.dtype)
+    memo = MemoizeOperator(op, max_neval=2)
+    first = op.H if adjoint_first else op
+    memo_first = memo.H if adjoint_first else memo
+    second = op if adjoint_first else op.H
+    memo_second = memo if adjoint_first else memo.H
+    vector = np.arange(1, first.shape[1] + 1, dtype=matrix.dtype)
+    intermediate = memo_first @ vector
+    expected = second @ (first @ vector)
+
+    result = memo_second @ intermediate
+
+    assert_array_almost_equal(result, expected, decimal=12)
+    assert memo.neval == 2
+    assert_array_almost_equal(memo_second @ intermediate, expected, decimal=12)
+    assert memo.neval == 2
+    assert_array_almost_equal(memo_first @ vector, first @ vector, decimal=12)
+    assert memo.neval == 2
+
+
+@pytest.mark.parametrize("max_neval", [1, 2, np.inf])
+@pytest.mark.parametrize("complex_values", [False, True])
+def test_memoize_cgls(max_neval, complex_values):
+    """Memoizing a non-unitary operator must preserve a least-squares solve."""
+    diagonal = np.array([2.0, 3.0, 4.0])
+    if complex_values:
+        diagonal = diagonal + 1j * np.array([1.0, -0.5, 2.0])
+    op = Diagonal(diagonal)
+    truth = np.array([1.0, -2.0, 0.5], dtype=diagonal.dtype)
+    data = op @ truth
+    memo = MemoizeOperator(op, max_neval=max_neval)
+
+    result = cgls(memo, data, np.zeros_like(truth), niter=3, tol=1e-14)[0]
+
+    assert_array_almost_equal(result, truth, decimal=12)
+    assert len(memo.store) <= max_neval
+    assert memo.neval > 1
+
+
+def test_memoize_direction_eviction():
+    """The cache budget remains shared across forward and adjoint evaluations."""
+    op = Diagonal(np.array([2.0, 3.0]))
+    memo = MemoizeOperator(op, max_neval=2)
+    x = np.array([1.0, 2.0])
+    y = op @ x
+    memo @ x
+    memo.H @ y
+    memo @ (2 * x)
+    assert memo.neval == 3
+    assert len(memo.store) == 2
+    assert_array_almost_equal(memo.H @ y, op.H @ y)
+    assert memo.neval == 3
+    assert_array_almost_equal(memo @ x, op @ x)
+    assert memo.neval == 4
+    assert len(memo.store) == 2
+
+
+@pytest.mark.parametrize("adjoint", [False, True])
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_memoize_nearby_inputs(adjoint, dtype):
+    """Nearby vectors are distinct evaluations, even within allclose tolerance."""
+    op = Diagonal(np.array([2.0, 3.0, 4.0], dtype=dtype))
+    memo = MemoizeOperator(op)
+    operator = op.H if adjoint else op
+    cached = memo.H if adjoint else memo
+    x = np.array([1.0, 2.0, 3.0], dtype=dtype)
+    nearby = x.copy()
+    nearby[0] += 1e-6
+    assert np.allclose(x, nearby)
+    cached @ x
+
+    result = cached @ nearby
+
+    np.testing.assert_array_equal(result, operator @ nearby)
+    assert memo.neval == 2
+    np.testing.assert_array_equal(cached @ nearby.copy(), result)
+    assert memo.neval == 2
+
+
+@pytest.mark.parametrize("adjoint", [False, True])
+@pytest.mark.parametrize("complex_values", [False, True])
+def test_memoize_input_dtype(adjoint, complex_values):
+    """Equal values with different precision must preserve output promotion."""
+    low = np.complex64 if complex_values else np.float32
+    high = np.complex128 if complex_values else np.float64
+    matrix = np.array([[2.0, 1.0], [0.0, 3.0]], dtype=low)
+    if complex_values:
+        matrix += 1j
+    op = MatrixMult(matrix, dtype=low)
+    memo = MemoizeOperator(op)
+    operator = op.H if adjoint else op
+    cached = memo.H if adjoint else memo
+    x = np.array([1.0, 2.0], dtype=low)
+    cached @ x
+    promoted = x.astype(high)
+
+    result = cached @ promoted
+    expected = operator @ promoted
+
+    assert result.dtype == expected.dtype
+    np.testing.assert_array_equal(result, expected)
+    assert memo.neval == 2