Skip to content

Commit 68986b0

Browse files
mrava87mrava87
authored
Merge pull request #811 from raashish1601/fix/803-scalar-division
Fix: scale the operator when dividing it by a scalar
2 parents 15f985b + f386b51 commit 68986b0

6 files changed

Lines changed: 93 additions & 15 deletions

File tree

‎pylops/linearoperator.py‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -835,14 +835,15 @@ def dot(self, x: NDArray) -> NDArray:
835835

836836
def div(
837837
self,
838-
y: NDArray,
838+
y: NDArray | float,
839839
niter: int = 100,
840840
densesolver: str = "scipy",
841-
) -> NDArray:
841+
) -> NDArray | LinearOperator:
842842
r"""Solve the linear problem :math:`\mathbf{y}=\mathbf{A}\mathbf{x}`.
843843
844844
Overloading of operator ``/`` to improve expressivity of `Pylops`
845-
when solving inverse problems.
845+
when solving inverse problems. Dividing by a scalar instead scales
846+
the operator, so ``Op / a`` is the same as ``Op * (1 / a)``.
846847
847848
Parameters
848849
----------
@@ -864,10 +865,13 @@ def div(
864865

865866
def __truediv__(
866867
self,
867-
y: NDArray,
868+
y: NDArray | float,
868869
niter: int = 100,
869870
densesolver: str = "scipy",
870-
) -> NDArray:
871+
) -> NDArray | LinearOperator:
872+
if np.isscalar(y):
873+
# A scalar divisor scales the operator, symmetric with ``Op * a``
874+
return self * (1 / y)
871875
if self.explicit is True:
872876
if sp.sparse.issparse(self.A):
873877
# use scipy solver for sparse matrices

‎pylops/signalprocessing/fft.py‎

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,10 @@ def _rmatvec(self, x: NDArray) -> NDArray:
140140
y = y.astype(self.rdtype)
141141
return y
142142

143-
def __truediv__(self, y: NDArray) -> NDArray:
143+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
144+
if np.isscalar(y):
145+
# A scalar divisor scales the operator (see LinearOperator.div)
146+
return super().__truediv__(y)
144147
if self.norm is not _FFTNorms.ORTHO:
145148
return self._rmatvec(y) / self._scale
146149
return self._rmatvec(y)
@@ -238,7 +241,10 @@ def _rmatvec(self, x: NDArray) -> NDArray:
238241
y = scipy.fft.fftshift(y, axes=self.axis)
239242
return y
240243

241-
def __truediv__(self, y: NDArray) -> NDArray:
244+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
245+
if np.isscalar(y):
246+
# A scalar divisor scales the operator (see LinearOperator.div)
247+
return super().__truediv__(y)
242248
if self.norm is not _FFTNorms.ORTHO:
243249
return self._rmatvec(y) / self._scale
244250
return self._rmatvec(y)
@@ -403,7 +409,10 @@ def _rmatvec(self, x: NDArray) -> NDArray:
403409
y = np.real(y)
404410
return y
405411

406-
def __truediv__(self, y: NDArray) -> NDArray:
412+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
413+
if np.isscalar(y):
414+
# A scalar divisor scales the operator (see LinearOperator.div)
415+
return super().__truediv__(y)
407416
if self.norm is _FFTNorms.ORTHO:
408417
return self._rmatvec(y)
409418
return self._rmatvec(y) / self._scale
@@ -495,7 +504,10 @@ def _rmatvec(self, x: NDArray) -> NDArray:
495504
y = scipy.fft.fftshift(y, axes=self.axis)
496505
return y
497506

498-
def __truediv__(self, y: NDArray) -> NDArray:
507+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
508+
if np.isscalar(y):
509+
# A scalar divisor scales the operator (see LinearOperator.div)
510+
return super().__truediv__(y)
499511
if self.norm is not _FFTNorms.ORTHO:
500512
return self._rmatvec(y) / self._scale
501513
return self._rmatvec(y)

‎pylops/signalprocessing/fft2d.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,10 @@ def _rmatvec(self, x):
136136
y = ncp.fft.fftshift(y, axes=self.axes[self.ifftshift_before])
137137
return y
138138

139-
def __truediv__(self, y: NDArray) -> NDArray:
139+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
140+
if np.isscalar(y):
141+
# A scalar divisor scales the operator (see LinearOperator.div)
142+
return super().__truediv__(y)
140143
if self.norm is not _FFTNorms.ORTHO:
141144
return self._rmatvec(y) / self._scale
142145
return self._rmatvec(y)
@@ -243,7 +246,10 @@ def _rmatvec(self, x):
243246
y = scipy.fft.fftshift(y, axes=self.axes[self.ifftshift_before])
244247
return y
245248

246-
def __truediv__(self, y: NDArray) -> NDArray:
249+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
250+
if np.isscalar(y):
251+
# A scalar divisor scales the operator (see LinearOperator.div)
252+
return super().__truediv__(y)
247253
if self.norm is not _FFTNorms.ORTHO:
248254
return self._rmatvec(y) / self._scale / self._scale
249255
return self._rmatvec(y)
@@ -357,7 +363,10 @@ def _rmatvec(self, x):
357363
y = scipy.fft.fftshift(y, axes=self.axes[self.ifftshift_before])
358364
return y
359365

360-
def __truediv__(self, y: NDArray) -> NDArray:
366+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
367+
if np.isscalar(y):
368+
# A scalar divisor scales the operator (see LinearOperator.div)
369+
return super().__truediv__(y)
361370
if self.norm is not _FFTNorms.ORTHO:
362371
return self._rmatvec(y) / self._scale / self._scale
363372
return self._rmatvec(y)

‎pylops/signalprocessing/fftnd.py‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,7 +122,10 @@ def _rmatvec(self, x: NDArray) -> NDArray:
122122
y = ncp.fft.fftshift(y, axes=self.axes[self.ifftshift_before])
123123
return y
124124

125-
def __truediv__(self, y: NDArray) -> NDArray:
125+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
126+
if np.isscalar(y):
127+
# A scalar divisor scales the operator (see LinearOperator.div)
128+
return super().__truediv__(y)
126129
if self.norm is not _FFTNorms.ORTHO:
127130
return self._rmatvec(y) / self._scale
128131
return self._rmatvec(y)
@@ -220,7 +223,10 @@ def _rmatvec(self, x: NDArray) -> NDArray:
220223
y = sp_fft.fftshift(y, axes=self.axes[self.ifftshift_before])
221224
return y
222225

223-
def __truediv__(self, y: NDArray) -> NDArray:
226+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
227+
if np.isscalar(y):
228+
# A scalar divisor scales the operator (see LinearOperator.div)
229+
return super().__truediv__(y)
224230
if self.norm is not _FFTNorms.ORTHO:
225231
return self._rmatvec(y) / self._scale
226232
return self._rmatvec(y)
@@ -320,7 +326,10 @@ def _rmatvec(self, x: NDArray) -> NDArray:
320326
y = scipy.fft.fftshift(y, axes=self.axes[self.ifftshift_before])
321327
return y
322328

323-
def __truediv__(self, y: NDArray) -> NDArray:
329+
def __truediv__(self, y: NDArray | float) -> NDArray | LinearOperator:
330+
if np.isscalar(y):
331+
# A scalar divisor scales the operator (see LinearOperator.div)
332+
return super().__truediv__(y)
324333
if self.norm is not _FFTNorms.ORTHO:
325334
return self._rmatvec(y) / self._scale
326335
return self._rmatvec(y)

‎pytests/test_ffts.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,32 @@ def test_unknown_engine(par):
318318
)
319319

320320

321+
@pytest.mark.parametrize("engine", ["numpy", "scipy", "fftw", "mkl_fft"])
322+
def test_scalar_division(engine):
323+
"""Check that dividing by a scalar scales the operator, like ``Op * (1 / a)``,
324+
instead of applying its inverse to the scalar
325+
"""
326+
if engine == "mkl_fft" and not mkl_fft_enabled:
327+
pytest.skip("mkl_fft is not installed")
328+
if engine == "fftw" and backend == "cupy":
329+
pytest.skip("fftw does not work with CuPy arrays")
330+
if engine == "scipy" and backend == "cupy":
331+
pytest.skip("scipy does not work with CuPy arrays")
332+
Ops = [FFT(dims=(8,), nfft=8, engine=engine)]
333+
if engine != "fftw":
334+
# FFT2D and FFTND have no fftw engine
335+
Ops += [
336+
FFT2D(dims=(4, 6), engine=engine),
337+
FFTND(dims=(2, 4, 6), engine=engine),
338+
]
339+
for Op in Ops:
340+
x = np.arange(Op.shape[1], dtype=np.float64)
341+
for scalar in (2, 2.0):
342+
Sop = Op / scalar
343+
assert Sop.shape == Op.shape
344+
assert_array_almost_equal(Sop @ x, (Op @ x) / 2, decimal=10)
345+
346+
321347
dtype_precision = [
322348
(np.float16, 1),
323349
(np.float32, 4),

‎pytests/test_linearoperator.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,24 @@ def test_scaled(par):
8181
assert S3op.dtype == dtype
8282

8383

84+
@pytest.mark.parametrize("par", [(par1), (par2), (par1j)])
85+
def test_scalar_division(par):
86+
"""Verify that dividing by a scalar scales the operator, like ``Op * (1 / a)``,
87+
instead of solving ``a = Op x``
88+
"""
89+
shape = (par["ny"], par["nx"])
90+
A = np.random.normal(0, 10, shape) + par["imag"] * np.random.normal(0, 10, shape)
91+
x = np.ones(par["nx"]) + par["imag"] * np.ones(par["nx"])
92+
Mop = MatrixMult(A, dtype=par["dtype"])
93+
# an explicit operator and one that is not
94+
for Op in (Mop, VStack([Mop])):
95+
for scalar in (2, 2.0):
96+
Sop = Op / scalar
97+
assert isinstance(Sop, LinearOperator)
98+
assert Sop.shape == Op.shape
99+
assert_array_almost_equal(Sop @ x, (Op @ x) / 2, decimal=6)
100+
101+
84102
@pytest.mark.parametrize("par", [(par1), (par1j)])
85103
def test_scipyop(par):
86104
"""Verify interaction between pylops and scipy Linear operators"""

0 commit comments

Comments
 (0)