Skip to content

Commit 9523ae1

Browse files
mrava87mrava87
committed
fix: ensure sliding ops work with fp32
1 parent 99f5aa8 commit 9523ae1

File tree

4 files changed

+250
-6
lines changed

4 files changed

+250
-6
lines changed

‎pylops/signalprocessing/sliding1d.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ def Sliding1D(
157157

158158
# create tapers
159159
if tapertype is not None:
160-
tap = taper(nwin, nover, tapertype=tapertype)
160+
tap = taper(nwin, nover, tapertype=tapertype).astype(Op.dtype)
161161
tapin = tap.copy()
162162
tapin[:nover] = 1
163163
tapend = tap.copy()
@@ -172,7 +172,9 @@ def Sliding1D(
172172
if tapertype is None:
173173
OOp = BlockDiag([Op for _ in range(nwins)])
174174
else:
175-
OOp = BlockDiag([Diagonal(taps[itap].ravel()) * Op for itap in range(nwins)])
175+
OOp = BlockDiag(
176+
[Diagonal(taps[itap].ravel(), dtype=Op.dtype) * Op for itap in range(nwins)]
177+
)
176178

177179
combining = HStack(
178180
[

‎pylops/signalprocessing/sliding1dNEW.py‎

Lines changed: 237 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,237 @@
1+
__all__ = [
2+
"sliding1d_design",
3+
"Sliding1D",
4+
]
5+
6+
import logging
7+
from typing import Tuple, Union
8+
9+
import numpy as np
10+
from numpy.lib.stride_tricks import sliding_window_view
11+
12+
from pylops import LinearOperator
13+
from pylops.signalprocessing.sliding2d import _slidingsteps
14+
from pylops.utils._internal import _value_or_sized_to_tuple
15+
from pylops.utils.backend import (
16+
get_array_module,
17+
get_sliding_window_view,
18+
to_cupy_conditional,
19+
)
20+
from pylops.utils.decorators import reshaped
21+
from pylops.utils.tapers import taper
22+
from pylops.utils.typing import InputDimsLike, NDArray
23+
24+
logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING)
25+
26+
27+
def sliding1d_design(
28+
dimd: int,
29+
nwin: int,
30+
nover: int,
31+
nop: int,
32+
) -> Tuple[int, int, Tuple[NDArray, NDArray], Tuple[NDArray, NDArray]]:
33+
"""Design Sliding1D operator
34+
35+
This routine can be used prior to creating the :class:`pylops.signalprocessing.Sliding1D`
36+
operator to identify the correct number of windows to be used based on the dimension of the data (``dimsd``),
37+
dimension of the window (``nwin``), overlap (``nover``),a and dimension of the operator acting in the model
38+
space.
39+
40+
Parameters
41+
----------
42+
dimsd : :obj:`tuple`
43+
Shape of 2-dimensional data.
44+
nwin : :obj:`tuple`
45+
Number of samples of window.
46+
nover : :obj:`tuple`
47+
Number of samples of overlapping part of window.
48+
nop : :obj:`tuple`
49+
Size of model in the transformed domain.
50+
51+
Returns
52+
-------
53+
nwins : :obj:`int`
54+
Number of windows.
55+
dim : :obj:`int`
56+
Shape of 2-dimensional model.
57+
mwins_inends : :obj:`tuple`
58+
Start and end indices for model patches.
59+
dwins_inends : :obj:`tuple`
60+
Start and end indices for data patches.
61+
62+
"""
63+
# data windows
64+
dwin_ins, dwin_ends = _slidingsteps(dimd, nwin, nover)
65+
dwins_inends = (dwin_ins, dwin_ends)
66+
nwins = len(dwin_ins)
67+
68+
# model windows
69+
dim = nwins * nop
70+
mwin_ins, mwin_ends = _slidingsteps(dim, nop, 0)
71+
mwins_inends = (mwin_ins, mwin_ends)
72+
73+
# print information about patching
74+
logging.warning("%d windows required...", nwins)
75+
logging.warning(
76+
"data wins - start:%s, end:%s",
77+
dwin_ins,
78+
dwin_ends,
79+
)
80+
logging.warning(
81+
"model wins - start:%s, end:%s",
82+
mwin_ins,
83+
mwin_ends,
84+
)
85+
return nwins, dim, mwins_inends, dwins_inends
86+
87+
88+
class Sliding1D(LinearOperator):
89+
r"""1D Sliding transform operator.
90+
91+
Apply a transform operator ``Op`` repeatedly to slices of the model
92+
vector in forward mode and slices of the data vector in adjoint mode.
93+
More specifically, in forward mode the model vector is divided into
94+
slices, each slice is transformed, and slices are then recombined in a
95+
sliding window fashion.
96+
97+
This operator can be used to perform local, overlapping transforms (e.g.,
98+
:obj:`pylops.signalprocessing.FFT`) on 1-dimensional arrays.
99+
100+
.. note:: The shape of the model has to be consistent with
101+
the number of windows for this operator not to return an error. As the
102+
number of windows depends directly on the choice of ``nwin`` and
103+
``nover``, it is recommended to first run ``sliding1d_design`` to obtain
104+
the corresponding ``dims`` and number of windows.
105+
106+
.. warning:: Depending on the choice of `nwin` and `nover` as well as the
107+
size of the data, sliding windows may not cover the entire data.
108+
The start and end indices of each window will be displayed and returned
109+
with running ``sliding1d_design``.
110+
111+
Parameters
112+
----------
113+
Op : :obj:`pylops.LinearOperator`
114+
Transform operator
115+
dim : :obj:`tuple`
116+
Shape of 1-dimensional model.
117+
dimd : :obj:`tuple`
118+
Shape of 1-dimensional data
119+
nwin : :obj:`int`
120+
Number of samples of window
121+
nover : :obj:`int`
122+
Number of samples of overlapping part of window
123+
tapertype : :obj:`str`, optional
124+
Type of taper (``hanning``, ``cosine``, ``cosinesquare`` or ``None``)
125+
name : :obj:`str`, optional
126+
.. versionadded:: 2.0.0
127+
128+
Name of operator (to be used by :func:`pylops.utils.describe.describe`)
129+
130+
Raises
131+
------
132+
ValueError
133+
Identified number of windows is not consistent with provided model
134+
shape (``dims``).
135+
136+
"""
137+
138+
def __init__(
139+
self,
140+
Op: LinearOperator,
141+
dim: Union[int, InputDimsLike],
142+
dimd: Union[int, InputDimsLike],
143+
nwin: int,
144+
nover: int,
145+
tapertype: str = "hanning",
146+
name: str = "S",
147+
) -> None:
148+
149+
dim: Tuple[int, ...] = _value_or_sized_to_tuple(dim)
150+
dimd: Tuple[int, ...] = _value_or_sized_to_tuple(dimd)
151+
152+
# data windows
153+
dwin_ins, dwin_ends = _slidingsteps(dimd[0], nwin, nover)
154+
self.dwin_inends = (dwin_ins, dwin_ends)
155+
nwins = len(dwin_ins)
156+
self.nwin = nwin
157+
self.nover = nover
158+
159+
# check windows
160+
if nwins * Op.shape[1] != dim[0] and Op.shape[1] != dim[0]:
161+
raise ValueError(
162+
f"Model shape (dim={dim}) is not consistent with chosen "
163+
f"number of windows. Run sliding1d_design to identify the "
164+
f"correct number of windows for the current "
165+
"model size..."
166+
)
167+
168+
# create tapers
169+
self.tapertype = tapertype
170+
if self.tapertype is not None:
171+
tap = taper(nwin, nover, tapertype=self.tapertype)
172+
tapin = tap.copy()
173+
tapin[:nover] = 1
174+
tapend = tap.copy()
175+
tapend[-nover:] = 1
176+
self.taps = [
177+
tapin,
178+
]
179+
for i in range(1, nwins - 1):
180+
self.taps.append(tap)
181+
self.taps.append(tapend)
182+
self.taps = np.vstack(self.taps)
183+
184+
# check if operator is applied to all windows simultaneously
185+
self.simOp = False
186+
if Op.shape[1] == dim[0]:
187+
self.simOp = True
188+
self.Op = Op
189+
190+
# create temporary shape and strides for cpy
191+
self.shape_wins = None
192+
self.strides_wins = None
193+
194+
super().__init__(
195+
dtype=Op.dtype,
196+
dims=(nwins, int(dim[0] // nwins)),
197+
dimsd=dimd,
198+
clinear=False,
199+
name=name,
200+
)
201+
202+
@reshaped
203+
def _matvec(self, x: NDArray) -> NDArray:
204+
ncp = get_array_module(x)
205+
if self.tapertype is not None:
206+
self.taps = to_cupy_conditional(x, self.taps)
207+
y = ncp.zeros(self.dimsd, dtype=self.dtype)
208+
if self.simOp:
209+
x = self.Op @ x
210+
for iwin0 in range(self.dims[0]):
211+
if self.simOp:
212+
xx = x[iwin0]
213+
else:
214+
xx = self.Op.matvec(x[iwin0])
215+
if self.tapertype is not None:
216+
xxwin = self.taps[iwin0] * xx
217+
else:
218+
xxwin = xx
219+
y[self.dwin_inends[0][iwin0] : self.dwin_inends[1][iwin0]] += xxwin
220+
return y
221+
222+
@reshaped
223+
def _rmatvec(self, x: NDArray) -> NDArray:
224+
ncp = get_array_module(x)
225+
ncp_sliding_window_view = get_sliding_window_view(x)
226+
if self.tapertype is not None:
227+
self.taps = to_cupy_conditional(x, self.taps)
228+
ywins = ncp_sliding_window_view(x, self.nwin)[:: self.nwin - self.nover]
229+
if self.tapertype is not None:
230+
ywins = ywins * self.taps
231+
if self.simOp:
232+
y = self.Op.H @ ywins
233+
else:
234+
y = ncp.zeros(self.dims, dtype=self.dtype)
235+
for iwin0 in range(self.dims[0]):
236+
y[iwin0] = self.Op.rmatvec(ywins[iwin0])
237+
return y

‎pylops/signalprocessing/sliding2d.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -191,7 +191,7 @@ def Sliding2D(
191191

192192
# create tapers
193193
if tapertype is not None:
194-
tap = taper2d(dimsd[1], nwin, nover, tapertype=tapertype)
194+
tap = taper2d(dimsd[1], nwin, nover, tapertype=tapertype).astype(Op.dtype)
195195
tapin = tap.copy()
196196
tapin[:nover] = 1
197197
tapend = tap.copy()
@@ -206,7 +206,9 @@ def Sliding2D(
206206
if tapertype is None:
207207
OOp = BlockDiag([Op for _ in range(nwins)])
208208
else:
209-
OOp = BlockDiag([Diagonal(taps[itap].ravel()) * Op for itap in range(nwins)])
209+
OOp = BlockDiag(
210+
[Diagonal(taps[itap].ravel(), dtype=Op.dtype) * Op for itap in range(nwins)]
211+
)
210212

211213
combining = HStack(
212214
[

‎pylops/signalprocessing/sliding3d.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -183,13 +183,16 @@ def Sliding3D(
183183

184184
# create tapers
185185
if tapertype is not None:
186-
tap = taper3d(dimsd[2], nwin, nover, tapertype=tapertype)
186+
tap = taper3d(dimsd[2], nwin, nover, tapertype=tapertype).astype(Op.dtype)
187187

188188
# transform to apply
189189
if tapertype is None:
190190
OOp = BlockDiag([Op for _ in range(nwins)], nproc=nproc)
191191
else:
192-
OOp = BlockDiag([Diagonal(tap.ravel()) * Op for _ in range(nwins)], nproc=nproc)
192+
OOp = BlockDiag(
193+
[Diagonal(tap.ravel(), dtype=Op.dtype) * Op for _ in range(nwins)],
194+
nproc=nproc,
195+
)
193196

194197
hstack = HStack(
195198
[

0 commit comments

Comments
 (0)