Skip to content

Commit dce8534

Browse files
mrava87mrava87
authored
Merge pull request #346 from cako/patch-swapaxes
Patch swapaxes
2 parents af3cf70 + ec101c5 commit dce8534

File tree

5 files changed

+70
-100
lines changed

5 files changed

+70
-100
lines changed

‎pylops/basicoperators/CausalIntegration.py‎

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -113,25 +113,22 @@ def __init__(
113113

114114
def _matvec(self, x):
115115
x = np.reshape(x, self.dims)
116-
if self.axis != -1:
117-
x = np.swapaxes(x, self.axis, -1)
116+
x = np.swapaxes(x, self.axis, -1)
118117
y = self.sampling * np.cumsum(x, axis=-1)
119118
if self.kind in ("half", "trapezoidal"):
120119
y -= self.sampling * x / 2.0
121120
if self.kind == "trapezoidal":
122121
y[..., 1:] -= self.sampling * x[..., 0:1] / 2.0
123122
if self.removefirst:
124123
y = y[..., 1:]
125-
if self.axis != -1:
126-
y = np.swapaxes(y, -1, self.axis)
124+
y = np.swapaxes(y, -1, self.axis)
127125
return y.ravel()
128126

129127
def _rmatvec(self, x):
130128
x = np.reshape(x, self.dimsd)
131129
if self.removefirst:
132130
x = np.insert(x, 0, 0, axis=self.axis)
133-
if self.axis != -1:
134-
x = np.swapaxes(x, self.axis, -1)
131+
x = np.swapaxes(x, self.axis, -1)
135132
xflip = np.flip(x, axis=-1)
136133
if self.kind == "half":
137134
y = self.sampling * (np.cumsum(xflip, axis=-1) - xflip / 2.0)
@@ -141,6 +138,5 @@ def _rmatvec(self, x):
141138
else:
142139
y = self.sampling * np.cumsum(xflip, axis=-1)
143140
y = np.flip(y, axis=-1)
144-
if self.axis != -1:
145-
y = np.swapaxes(y, -1, self.axis)
141+
y = np.swapaxes(y, -1, self.axis)
146142
return y.ravel()

‎pylops/basicoperators/FirstDerivative.py‎

Lines changed: 27 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -99,86 +99,74 @@ def __init__(
9999
def _matvec_forward(self, x):
100100
ncp = get_array_module(x)
101101
x = ncp.reshape(x, self.dims)
102-
if self.axis > 0: # need to bring the dim. to derive to first dim.
103-
x = ncp.swapaxes(x, self.axis, 0)
102+
x = ncp.swapaxes(x, self.axis, -1)
104103
y = ncp.zeros(x.shape, self.dtype)
105-
y[:-1, ...] = (x[1:, ...] - x[:-1, ...]) / self.sampling
106-
if self.axis > 0:
107-
y = ncp.swapaxes(y, 0, self.axis)
104+
y[..., :-1] = (x[..., 1:] - x[..., :-1]) / self.sampling
105+
y = ncp.swapaxes(y, -1, self.axis)
108106
y = y.ravel()
109107
return y
110108

111109
def _rmatvec_forward(self, x):
112110
ncp = get_array_module(x)
113111
x = ncp.reshape(x, self.dims)
114-
if self.axis > 0: # need to bring the dim. to derive to first dim.
115-
x = ncp.swapaxes(x, self.axis, 0)
112+
x = ncp.swapaxes(x, self.axis, -1)
116113
y = ncp.zeros(x.shape, self.dtype)
117-
y[:-1, ...] -= x[:-1, ...]
118-
y[1:, ...] += x[:-1, ...]
114+
y[..., :-1] -= x[..., :-1]
115+
y[..., 1:] += x[..., :-1]
119116
y /= self.sampling
120-
if self.axis > 0:
121-
y = ncp.swapaxes(y, 0, self.axis)
117+
y = ncp.swapaxes(y, -1, self.axis)
122118
y = y.ravel()
123119
return y
124120

125121
def _matvec_centered(self, x):
126122
ncp = get_array_module(x)
127123
x = ncp.reshape(x, self.dims)
128-
if self.axis > 0: # need to bring the dim. to derive to first dim.
129-
x = ncp.swapaxes(x, self.axis, 0)
124+
x = ncp.swapaxes(x, self.axis, -1)
130125
y = ncp.zeros(x.shape, self.dtype)
131-
y[1:-1, ...] = 0.5 * x[2:, ...] - 0.5 * x[0:-2, ...]
126+
y[..., 1:-1] = 0.5 * (x[..., 2:] - x[..., :-2])
132127
if self.edge:
133-
y[0, ...] = x[1, ...] - x[0, ...]
134-
y[-1, ...] = x[-1, ...] - x[-2, ...]
128+
y[..., 0] = x[..., 1] - x[..., 0]
129+
y[..., -1] = x[..., -1] - x[..., -2]
135130
y /= self.sampling
136-
if self.axis > 0:
137-
y = ncp.swapaxes(y, 0, self.axis)
131+
y = ncp.swapaxes(y, -1, self.axis)
138132
y = y.ravel()
139133
return y
140134

141135
def _rmatvec_centered(self, x):
142136
ncp = get_array_module(x)
143137
x = ncp.reshape(x, self.dims)
144-
if self.axis > 0: # need to bring the dim. to derive to first dim.
145-
x = ncp.swapaxes(x, self.axis, 0)
138+
x = ncp.swapaxes(x, self.axis, -1)
146139
y = ncp.zeros(x.shape, self.dtype)
147-
y[0:-2, ...] -= 0.5 * x[1:-1, ...]
148-
y[2:, ...] += 0.5 * x[1:-1, ...]
140+
y[..., :-2] -= 0.5 * x[..., 1:-1]
141+
y[..., 2:] += 0.5 * x[..., 1:-1]
149142
if self.edge:
150-
y[0, ...] -= x[0, ...]
151-
y[1, ...] += x[0, ...]
152-
y[-2, ...] -= x[-1, ...]
153-
y[-1, ...] += x[-1, ...]
143+
y[..., 0] -= x[..., 0]
144+
y[..., 1] += x[..., 0]
145+
y[..., -2] -= x[..., -1]
146+
y[..., -1] += x[..., -1]
154147
y /= self.sampling
155-
if self.axis > 0:
156-
y = ncp.swapaxes(y, 0, self.axis)
148+
y = ncp.swapaxes(y, -1, self.axis)
157149
y = y.ravel()
158150
return y
159151

160152
def _matvec_backward(self, x):
161153
ncp = get_array_module(x)
162154
x = ncp.reshape(x, self.dims)
163-
if self.axis > 0: # need to bring the dim. to derive to first dim.
164-
x = ncp.swapaxes(x, self.axis, 0)
155+
x = ncp.swapaxes(x, self.axis, -1)
165156
y = ncp.zeros(x.shape, self.dtype)
166-
y[1:, ...] = (x[1:, ...] - x[:-1, ...]) / self.sampling
167-
if self.axis > 0:
168-
y = ncp.swapaxes(y, 0, self.axis)
157+
y[..., 1:] = (x[..., 1:] - x[..., :-1]) / self.sampling
158+
y = ncp.swapaxes(y, -1, self.axis)
169159
y = y.ravel()
170160
return y
171161

172162
def _rmatvec_backward(self, x):
173163
ncp = get_array_module(x)
174164
x = ncp.reshape(x, self.dims)
175-
if self.axis > 0: # need to bring the dim. to derive to first dim.
176-
x = ncp.swapaxes(x, self.axis, 0)
165+
x = ncp.swapaxes(x, self.axis, -1)
177166
y = ncp.zeros(x.shape, self.dtype)
178-
y[:-1, ...] -= x[1:, ...]
179-
y[1:, ...] += x[1:, ...]
167+
y[..., :-1] -= x[..., 1:]
168+
y[..., 1:] += x[..., 1:]
180169
y /= self.sampling
181-
if self.axis > 0:
182-
y = ncp.swapaxes(y, 0, self.axis)
170+
y = ncp.swapaxes(y, -1, self.axis)
183171
y = y.ravel()
184172
return y

‎pylops/basicoperators/Restriction.py‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -155,14 +155,12 @@ def mask(self, x):
155155

156156
y = np_ma.array(np.zeros(self.dims), mask=np.ones(self.dims), dtype=self.dtype)
157157
x = np.reshape(x, self.dims)
158-
if self.axis > 0:
159-
x = np.swapaxes(x, self.axis, 0)
160-
y = np.swapaxes(y, self.axis, 0)
161-
y.mask[iava] = False
158+
x = np.swapaxes(x, self.axis, -1)
159+
y = np.swapaxes(y, self.axis, -1)
160+
y.mask[..., iava] = False
162161
if ncp == np:
163-
y[iava] = x[self.iava]
162+
y[..., iava] = x[..., self.iava]
164163
else:
165-
y[iava] = ncp.asnumpy(x)[iava]
166-
if self.axis > 0:
167-
y = np.swapaxes(y, 0, self.axis)
164+
y[..., iava] = ncp.asnumpy(x)[..., iava]
165+
y = np.swapaxes(y, -1, self.axis)
168166
return y

‎pylops/basicoperators/SecondDerivative.py‎

Lines changed: 24 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -100,65 +100,57 @@ def __init__(
100100
def _matvec_forward(self, x):
101101
ncp = get_array_module(x)
102102
x = ncp.reshape(x, self.dims)
103-
if self.axis > 0: # need to bring the dim. to derive to first dim.
104-
x = ncp.swapaxes(x, self.axis, 0)
103+
x = ncp.swapaxes(x, self.axis, -1)
105104
y = ncp.zeros(x.shape, self.dtype)
106-
y[:-2, ...] = x[2:, ...] - 2 * x[1:-1, ...] + x[0:-2, ...]
105+
y[..., :-2] = x[..., 2:] - 2 * x[..., 1:-1] + x[..., :-2]
107106
y /= self.sampling ** 2
108-
if self.axis > 0:
109-
y = ncp.swapaxes(y, 0, self.axis)
107+
y = ncp.swapaxes(y, -1, self.axis)
110108
y = y.ravel()
111109
return y
112110

113111
def _rmatvec_forward(self, x):
114112
ncp = get_array_module(x)
115113
x = ncp.reshape(x, self.dims)
116-
if self.axis > 0: # need to bring the dim. to derive to first dim.
117-
x = ncp.swapaxes(x, self.axis, 0)
114+
x = ncp.swapaxes(x, self.axis, -1)
118115
y = ncp.zeros(x.shape, self.dtype)
119-
y[0:-2, ...] += x[:-2, ...]
120-
y[1:-1, ...] -= 2 * x[:-2, ...]
121-
y[2:, ...] += x[:-2, ...]
116+
y[..., :-2] += x[..., :-2]
117+
y[..., 1:-1] -= 2 * x[..., :-2]
118+
y[..., 2:] += x[..., :-2]
122119
y /= self.sampling ** 2
123-
if self.axis > 0:
124-
y = ncp.swapaxes(y, 0, self.axis)
120+
y = ncp.swapaxes(y, -1, self.axis)
125121
y = y.ravel()
126122
return y
127123

128124
def _matvec_centered(self, x):
129125
ncp = get_array_module(x)
130126
x = ncp.reshape(x, self.dims)
131-
if self.axis > 0: # need to bring the dim. to derive to first dim.
132-
x = ncp.swapaxes(x, self.axis, 0)
127+
x = ncp.swapaxes(x, self.axis, -1)
133128
y = ncp.zeros(x.shape, self.dtype)
134-
y[1:-1, ...] = x[2:, ...] - 2 * x[1:-1, ...] + x[0:-2, ...]
129+
y[..., 1:-1] = x[..., 2:] - 2 * x[..., 1:-1] + x[..., :-2]
135130
if self.edge:
136-
y[0, ...] = x[0, ...] - 2 * x[1, ...] + x[2, ...]
137-
y[-1, ...] = x[-3, ...] - 2 * x[-2, ...] + x[-1, ...]
131+
y[..., 0] = x[..., 0] - 2 * x[..., 1] + x[..., 2]
132+
y[..., -1] = x[..., -3] - 2 * x[..., -2] + x[..., -1]
138133
y /= self.sampling ** 2
139-
if self.axis > 0:
140-
y = ncp.swapaxes(y, 0, self.axis)
134+
y = ncp.swapaxes(y, -1, self.axis)
141135
y = y.ravel()
142136
return y
143137

144138
def _rmatvec_centered(self, x):
145139
ncp = get_array_module(x)
146140
x = ncp.reshape(x, self.dims)
147-
if self.axis > 0: # need to bring the dim. to derive to first dim.
148-
x = ncp.swapaxes(x, self.axis, 0)
141+
x = ncp.swapaxes(x, self.axis, -1)
149142
y = ncp.zeros(x.shape, self.dtype)
150-
y[0:-2, ...] += x[1:-1, ...]
151-
y[1:-1, ...] -= 2 * x[1:-1, ...]
152-
y[2:, ...] += x[1:-1, ...]
143+
y[..., :-2] += x[..., 1:-1]
144+
y[..., 1:-1] -= 2 * x[..., 1:-1]
145+
y[..., 2:] += x[..., 1:-1]
153146
if self.edge:
154-
y[0, ...] += x[0, ...]
155-
y[1, ...] -= 2 * x[0, ...]
156-
y[2, ...] += x[0, ...]
157-
y[-3, ...] += x[-1, ...]
158-
y[-2, ...] -= 2 * x[-1, ...]
159-
y[-1, ...] += x[-1, ...]
147+
y[..., 0] += x[..., 0]
148+
y[..., 1] -= 2 * x[..., 0]
149+
y[..., 2] += x[..., 0]
150+
y[..., -3] += x[..., -1]
151+
y[..., -2] -= 2 * x[..., -1]
152+
y[..., -1] += x[..., -1]
160153
y /= self.sampling ** 2
161-
if self.axis > 0:
162-
y = ncp.swapaxes(y, 0, self.axis)
154+
y = ncp.swapaxes(y, -1, self.axis)
163155
y = y.ravel()
164156
return y

‎pylops/basicoperators/Symmetrize.py‎

Lines changed: 9 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -74,24 +74,20 @@ def _matvec(self, x):
7474
ncp = get_array_module(x)
7575
y = ncp.zeros(self.dimsd, dtype=self.dtype)
7676
x = ncp.reshape(x, self.dims)
77-
if self.axis > 0: # bring the dimension to symmetrize to first
78-
x = ncp.swapaxes(x, self.axis, 0)
79-
y = ncp.swapaxes(y, self.axis, 0)
80-
y[self.nsym - 1 :] = x
81-
y[: self.nsym - 1] = x[-1:0:-1]
82-
if self.axis > 0:
83-
y = ncp.swapaxes(y, 0, self.axis)
77+
x = ncp.swapaxes(x, self.axis, -1)
78+
y = ncp.swapaxes(y, self.axis, -1)
79+
y[..., self.nsym - 1 :] = x
80+
y[..., : self.nsym - 1] = x[..., -1:0:-1]
81+
y = ncp.swapaxes(y, -1, self.axis)
8482
y = y.ravel()
8583
return y
8684

8785
def _rmatvec(self, x):
8886
ncp = get_array_module(x)
8987
x = ncp.reshape(x, self.dimsd)
90-
if self.axis > 0: # bring the dimension to symmetrize to first
91-
x = ncp.swapaxes(x, self.axis, 0)
92-
y = x[self.nsym - 1 :].copy()
93-
y[1:] += x[self.nsym - 2 :: -1]
94-
if self.axis > 0:
95-
y = ncp.swapaxes(y, 0, self.axis)
88+
x = ncp.swapaxes(x, self.axis, -1)
89+
y = x[..., self.nsym - 1 :].copy()
90+
y[..., 1:] += x[..., self.nsym - 2 :: -1]
91+
y = ncp.swapaxes(y, -1, self.axis)
9692
y = y.ravel()
9793
return y

0 commit comments

Comments
 (0)