diff --git a/examples/plot_memoizeoperator.py b/examples/plot_memoizeoperator.py new file mode 100644 index 00000000..d497b58c --- /dev/null +++ b/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) diff --git a/pylops/basicoperators/memoizeoperator.py b/pylops/basicoperators/memoizeoperator.py index 1ad4b4f4..f2a9c3f9 100644 --- a/pylops/basicoperators/memoizeoperator.py +++ b/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 diff --git a/pytests/test_memoizeoperator.py b/pytests/test_memoizeoperator.py index a92f725f..7eccb6d4 100644 --- a/pytests/test_memoizeoperator.py +++ b/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):