diff --git a/pylops/linearoperator.py b/pylops/linearoperator.py index 7a47ac0d..da8fcd1e 100644 --- a/pylops/linearoperator.py +++ b/pylops/linearoperator.py @@ -835,14 +835,15 @@ def dot(self, x: NDArray) -> NDArray: def div( self, - y: NDArray, + y: NDArray | float, niter: int = 100, densesolver: str = "scipy", - ) -> NDArray: + ) -> NDArray | LinearOperator: r"""Solve the linear problem :math:`\mathbf{y}=\mathbf{A}\mathbf{x}`. Overloading of operator ``/`` to improve expressivity of `Pylops` - when solving inverse problems. + when solving inverse problems. Dividing by a scalar instead scales + the operator, so ``Op / a`` is the same as ``Op * (1 / a)``. Parameters ---------- @@ -864,10 +865,13 @@ def div( def __truediv__( self, - y: NDArray, + y: NDArray | float, niter: int = 100, densesolver: str = "scipy", - ) -> NDArray: + ) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator, symmetric with ``Op * a`` + return self * (1 / y) if self.explicit is True: if sp.sparse.issparse(self.A): # use scipy solver for sparse matrices diff --git a/pylops/signalprocessing/fft.py b/pylops/signalprocessing/fft.py index fb234dcc..8322f75f 100644 --- a/pylops/signalprocessing/fft.py +++ b/pylops/signalprocessing/fft.py @@ -140,7 +140,10 @@ def _rmatvec(self, x: NDArray) -> NDArray: y = y.astype(self.rdtype) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale return self._rmatvec(y) @@ -238,7 +241,10 @@ def _rmatvec(self, x: NDArray) -> NDArray: y = scipy.fft.fftshift(y, axes=self.axis) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale return self._rmatvec(y) @@ -403,7 +409,10 @@ def _rmatvec(self, x: NDArray) -> NDArray: y = np.real(y) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is _FFTNorms.ORTHO: return self._rmatvec(y) return self._rmatvec(y) / self._scale @@ -495,7 +504,10 @@ def _rmatvec(self, x: NDArray) -> NDArray: y = scipy.fft.fftshift(y, axes=self.axis) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale return self._rmatvec(y) diff --git a/pylops/signalprocessing/fft2d.py b/pylops/signalprocessing/fft2d.py index e05d6d21..eb156159 100644 --- a/pylops/signalprocessing/fft2d.py +++ b/pylops/signalprocessing/fft2d.py @@ -136,7 +136,10 @@ def _rmatvec(self, x): y = ncp.fft.fftshift(y, axes=self.axes[self.ifftshift_before]) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale return self._rmatvec(y) @@ -243,7 +246,10 @@ def _rmatvec(self, x): y = scipy.fft.fftshift(y, axes=self.axes[self.ifftshift_before]) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale / self._scale return self._rmatvec(y) @@ -357,7 +363,10 @@ def _rmatvec(self, x): y = scipy.fft.fftshift(y, axes=self.axes[self.ifftshift_before]) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale / self._scale return self._rmatvec(y) diff --git a/pylops/signalprocessing/fftnd.py b/pylops/signalprocessing/fftnd.py index 980ffa11..54d28e55 100644 --- a/pylops/signalprocessing/fftnd.py +++ b/pylops/signalprocessing/fftnd.py @@ -122,7 +122,10 @@ def _rmatvec(self, x: NDArray) -> NDArray: y = ncp.fft.fftshift(y, axes=self.axes[self.ifftshift_before]) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale return self._rmatvec(y) @@ -220,7 +223,10 @@ def _rmatvec(self, x: NDArray) -> NDArray: y = sp_fft.fftshift(y, axes=self.axes[self.ifftshift_before]) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale return self._rmatvec(y) @@ -320,7 +326,10 @@ def _rmatvec(self, x: NDArray) -> NDArray: y = scipy.fft.fftshift(y, axes=self.axes[self.ifftshift_before]) return y - def __truediv__(self, y: NDArray) -> NDArray: + def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator: + if np.isscalar(y): + # A scalar divisor scales the operator (see LinearOperator.div) + return super().__truediv__(y) if self.norm is not _FFTNorms.ORTHO: return self._rmatvec(y) / self._scale return self._rmatvec(y) diff --git a/pytests/test_ffts.py b/pytests/test_ffts.py index 188ddaa9..ac93c6d4 100644 --- a/pytests/test_ffts.py +++ b/pytests/test_ffts.py @@ -318,6 +318,32 @@ def test_unknown_engine(par): ) +@pytest.mark.parametrize("engine", ["numpy", "scipy", "fftw", "mkl_fft"]) +def test_scalar_division(engine): + """Check that dividing by a scalar scales the operator, like ``Op * (1 / a)``, + instead of applying its inverse to the scalar + """ + if engine == "mkl_fft" and not mkl_fft_enabled: + pytest.skip("mkl_fft is not installed") + if engine == "fftw" and backend == "cupy": + pytest.skip("fftw does not work with CuPy arrays") + if engine == "scipy" and backend == "cupy": + pytest.skip("scipy does not work with CuPy arrays") + Ops = [FFT(dims=(8,), nfft=8, engine=engine)] + if engine != "fftw": + # FFT2D and FFTND have no fftw engine + Ops += [ + FFT2D(dims=(4, 6), engine=engine), + FFTND(dims=(2, 4, 6), engine=engine), + ] + for Op in Ops: + x = np.arange(Op.shape[1], dtype=np.float64) + for scalar in (2, 2.0): + Sop = Op / scalar + assert Sop.shape == Op.shape + assert_array_almost_equal(Sop @ x, (Op @ x) / 2, decimal=10) + + dtype_precision = [ (np.float16, 1), (np.float32, 4), diff --git a/pytests/test_linearoperator.py b/pytests/test_linearoperator.py index 29eba9c8..cca6b5d1 100644 --- a/pytests/test_linearoperator.py +++ b/pytests/test_linearoperator.py @@ -81,6 +81,24 @@ def test_scaled(par): assert S3op.dtype == dtype +@pytest.mark.parametrize("par", [(par1), (par2), (par1j)]) +def test_scalar_division(par): + """Verify that dividing by a scalar scales the operator, like ``Op * (1 / a)``, + instead of solving ``a = Op x`` + """ + shape = (par["ny"], par["nx"]) + A = np.random.normal(0, 10, shape) + par["imag"] * np.random.normal(0, 10, shape) + x = np.ones(par["nx"]) + par["imag"] * np.ones(par["nx"]) + Mop = MatrixMult(A, dtype=par["dtype"]) + # an explicit operator and one that is not + for Op in (Mop, VStack([Mop])): + for scalar in (2, 2.0): + Sop = Op / scalar + assert isinstance(Sop, LinearOperator) + assert Sop.shape == Op.shape + assert_array_almost_equal(Sop @ x, (Op @ x) / 2, decimal=6) + + @pytest.mark.parametrize("par", [(par1), (par1j)]) def test_scipyop(par): """Verify interaction between pylops and scipy Linear operators"""