From dec5089b0216d62c3dc2c55a4dadc6bb3279f0b4 Mon Sep 17 00:00:00 2001 From: AHMETHAKANBEZIR1 Date: Thu, 15 Oct 2026 23:09:48 +0300 Subject: [PATCH 1/2] Fix: distinguish forward and adjoint memoized evaluations --- examples/plot_memoizeoperator.py | 31 +++++++++++ pylops/basicoperators/memoizeoperator.py | 21 ++++++-- pytests/test_memoizeoperator.py | 68 +++++++++++++++++++++++- 3 files changed, 114 insertions(+), 6 deletions(-) create mode 100644 examples/plot_memoizeoperator.py diff --git a/examples/plot_memoizeoperator.py b/examples/plot_memoizeoperator.py new file mode 100644 index 00000000..77159cbd --- /dev/null +++ b/examples/plot_memoizeoperator.py @@ -0,0 +1,31 @@ +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 diff --git a/pylops/basicoperators/memoizeoperator.py b/pylops/basicoperators/memoizeoperator.py index 1ad4b4f4..c2b09117 100644 --- a/pylops/basicoperators/memoizeoperator.py +++ b/pylops/basicoperators/memoizeoperator.py @@ -13,7 +13,9 @@ 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. Parameters ---------- @@ -43,26 +45,35 @@ 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 np.allclose(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 np.allclose(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 diff --git a/pytests/test_memoizeoperator.py b/pytests/test_memoizeoperator.py index a92f725f..9f9963bd 100644 --- a/pytests/test_memoizeoperator.py +++ b/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,68 @@ 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 From 40e49a2301bcd5fe7cb3007393ebf6d3d1c4b74b Mon Sep 17 00:00:00 2001 From: AHMETHAKANBEZIR1 Date: Thu, 15 Oct 2026 23:21:50 +0300 Subject: [PATCH 2/2] Fix: require exact memoized input values and dtype --- examples/plot_memoizeoperator.py | 8 ++++ pylops/basicoperators/memoizeoperator.py | 11 ++++-- pytests/test_memoizeoperator.py | 47 ++++++++++++++++++++++++ 3 files changed, 63 insertions(+), 3 deletions(-) diff --git a/examples/plot_memoizeoperator.py b/examples/plot_memoizeoperator.py index 77159cbd..abd04125 100644 --- a/examples/plot_memoizeoperator.py +++ b/examples/plot_memoizeoperator.py @@ -29,3 +29,11 @@ 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 diff --git a/pylops/basicoperators/memoizeoperator.py b/pylops/basicoperators/memoizeoperator.py index c2b09117..940dd214 100644 --- a/pylops/basicoperators/memoizeoperator.py +++ b/pylops/basicoperators/memoizeoperator.py @@ -15,7 +15,8 @@ class MemoizeOperator(LinearOperator): and stores the last ``max_neval`` model/data 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. + the inverse of a forward evaluation. Cache hits require identical input + values and dtype; nearby vectors are evaluated independently. Parameters ---------- @@ -52,7 +53,11 @@ def _matvec(self, x: NDArray) -> NDArray: for (xstored, ystored), is_adjoint in zip( self.store, self._store_is_adjoint, strict=True ): - if not is_adjoint and np.allclose(xstored, x): + 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 @@ -67,7 +72,7 @@ def _rmatvec(self, y: NDArray) -> NDArray: for (xstored, ystored), is_adjoint in zip( self.store, self._store_is_adjoint, strict=True ): - if is_adjoint and np.allclose(ystored, y): + 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 diff --git a/pytests/test_memoizeoperator.py b/pytests/test_memoizeoperator.py index 9f9963bd..0b4a8c37 100644 --- a/pytests/test_memoizeoperator.py +++ b/pytests/test_memoizeoperator.py @@ -146,3 +146,50 @@ def test_memoize_direction_eviction(): 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