1 from __future__ import division
2
3 import numpy as num
4 import _correlate
5 import numpy.fft as dft
6 import iraf_frame
7
8 VALID = 0
9 SAME = 1
10 FULL = 2
11 PASS = 3
12
13 convolution_modes = {
14 "valid":0,
15 "same":1,
16 "full":2,
17 "pass":3,
18 }
19
31
33 """correlate(data, kernel, mode=FULL)
34
35 >>> correlate(num.arange(8), [1, 2], mode=VALID)
36 array([ 2, 5, 8, 11, 14, 17, 20])
37 >>> correlate(num.arange(8), [1, 2], mode=SAME)
38 array([ 0, 2, 5, 8, 11, 14, 17, 20])
39 >>> correlate(num.arange(8), [1, 2], mode=FULL)
40 array([ 0, 2, 5, 8, 11, 14, 17, 20, 7])
41 >>> correlate(num.arange(8), [1, 2, 3], mode=VALID)
42 array([ 8, 14, 20, 26, 32, 38])
43 >>> correlate(num.arange(8), [1, 2, 3], mode=SAME)
44 array([ 3, 8, 14, 20, 26, 32, 38, 20])
45 >>> correlate(num.arange(8), [1, 2, 3], mode=FULL)
46 array([ 0, 3, 8, 14, 20, 26, 32, 38, 20, 7])
47 >>> correlate(num.arange(8), [1, 2, 3, 4, 5, 6], mode=VALID)
48 array([ 70, 91, 112])
49 >>> correlate(num.arange(8), [1, 2, 3, 4, 5, 6], mode=SAME)
50 array([ 17, 32, 50, 70, 91, 112, 85, 60])
51 >>> correlate(num.arange(8), [1, 2, 3, 4, 5, 6], mode=FULL)
52 array([ 0, 6, 17, 32, 50, 70, 91, 112, 85, 60, 38, 20, 7])
53 >>> correlate(num.arange(8), 1+1j)
54 Traceback (most recent call last):
55 ...
56 TypeError: array cannot be safely cast to required type
57
58 """
59 data, kernel = _condition_inputs(data, kernel)
60 lenk = len(kernel)
61 halfk = int(lenk/2)
62 even = (lenk % 2 == 0)
63 kdata = [0] * lenk
64
65 if mode in convolution_modes.keys():
66 mode = convolution_modes[ mode ]
67
68 result_type = max(kernel.dtype.name, data.dtype.name)
69
70 if mode == VALID:
71 wdata = num.concatenate((kdata, data, kdata))
72 result = wdata.astype(result_type)
73 _correlate.Correlate1d(kernel, wdata, result)
74 return result[lenk+halfk:-lenk-halfk+even]
75 elif mode == SAME:
76 wdata = num.concatenate((kdata, data, kdata))
77 result = wdata.astype(result_type)
78 _correlate.Correlate1d(kernel, wdata, result)
79 return result[lenk:-lenk]
80 elif mode == FULL:
81 wdata = num.concatenate((kdata, data, kdata))
82 result = wdata.astype(result_type)
83 _correlate.Correlate1d(kernel, wdata, result)
84 return result[halfk+1:-halfk-1+even]
85 elif mode == PASS:
86 result = data.astype(result_type)
87 _correlate.Correlate1d(kernel, data, result)
88 return result
89 else:
90 raise ValueError("Invalid convolution mode.")
91
92 cross_correlate = correlate
93
94 pix_modes = {
95 "nearest" : 0,
96 "reflect": 1,
97 "wrap" : 2,
98 "constant": 3
99 }
100
102 """convolve(data, kernel, mode=FULL)
103 Returns the discrete, linear convolution of 1-D
104 sequences a and v; mode can be 0 (VALID), 1 (SAME), or 2 (FULL)
105 to specify size of the resulting sequence.
106
107 >>> convolve(num.arange(8), [1, 2], mode=VALID)
108 array([ 1, 4, 7, 10, 13, 16, 19])
109 >>> convolve(num.arange(8), [1, 2], mode=SAME)
110 array([ 0, 1, 4, 7, 10, 13, 16, 19])
111 >>> convolve(num.arange(8), [1, 2], mode=FULL)
112 array([ 0, 1, 4, 7, 10, 13, 16, 19, 14])
113 >>> convolve(num.arange(8), [1, 2, 3], mode=VALID)
114 array([ 4, 10, 16, 22, 28, 34])
115 >>> convolve(num.arange(8), [1, 2, 3], mode=SAME)
116 array([ 1, 4, 10, 16, 22, 28, 34, 32])
117 >>> convolve(num.arange(8), [1, 2, 3], mode=FULL)
118 array([ 0, 1, 4, 10, 16, 22, 28, 34, 32, 21])
119 >>> convolve(num.arange(8), [1, 2, 3, 4, 5, 6], mode=VALID)
120 array([35, 56, 77])
121 >>> convolve(num.arange(8), [1, 2, 3, 4, 5, 6], mode=SAME)
122 array([ 4, 10, 20, 35, 56, 77, 90, 94])
123 >>> convolve(num.arange(8), [1, 2, 3, 4, 5, 6], mode=FULL)
124 array([ 0, 1, 4, 10, 20, 35, 56, 77, 90, 94, 88, 71, 42])
125 >>> convolve([1.,2.], num.arange(10.))
126 array([ 0., 1., 4., 7., 10., 13., 16., 19., 22., 25., 18.])
127 """
128 data, kernel = _condition_inputs(data, kernel)
129 if len(data) >= len(kernel):
130 return correlate(data, kernel[::-1], mode)
131 else:
132 return correlate(kernel, data[::-1], mode)
133
134
136 ox = num.arange(mew-sigmas*sigma,
137 mew+sigmas*sigma,
138 2*sigmas*sigma/npoints, type=num.float64)
139 x = ox-mew
140 x /= sigma
141 x = x * x
142 x *= -1/2
143 x = num.exp(x)
144 return ox, 1/(sigma * num.sqrt(2*num.pi)) * x
145
147 """_correlate2d_fft does 2d correlation of 'data' with 'kernel', storing
148 the result in 'output' using the FFT to perform the correlation.
149
150 supported 'mode's include:
151 'nearest' elements beyond boundary come from nearest edge pixel.
152 'wrap' elements beyond boundary come from the opposite array edge.
153 'reflect' elements beyond boundary come from reflection on same array edge.
154 'constant' elements beyond boundary are set to 'cval'
155 """
156 shape = data0.shape
157 kshape = kernel0.shape
158 oversized = (num.array(shape) + num.array(kshape))
159
160 dy = kshape[0] // 2
161 dx = kshape[1] // 2
162
163 kernel = num.zeros(oversized, dtype=num.float64)
164 kernel[:kshape[0], :kshape[1]] = kernel0[::-1,::-1]
165 data = iraf_frame.frame(data0, oversized, mode=mode, cval=cval)
166
167 complex_result = (isinstance(data, num.complexfloating) or
168 isinstance(kernel, num.complexfloating))
169
170 Fdata = dft.fft2(data)
171 del data
172
173 Fkernel = dft.fft2(kernel)
174 del kernel
175
176 num.multiply(Fdata, Fkernel, Fdata)
177 del Fkernel
178
179 if complex_result:
180 convolved = dft.irfft2( Fdata, s=oversized)
181 else:
182 convolved = dft.irfft2( Fdata, s=oversized)
183
184 result = convolved[ kshape[0]-1:shape[0]+kshape[0]-1, kshape[1]-1:shape[1]+kshape[1]-1 ]
185
186 if output is not None:
187 output._copyFrom( result )
188 else:
189 return result
190
191
194
196 """The _correlate.Correlate2d C-code can only handle kernels which
197 fit inside the data array. Since convolution and correlation are
198 commutative, _fix_data_kernel reverses kernel and data if necessary
199 and panics if there's no good order.
200 """
201 data, kernel = map(num.asarray, [data, kernel])
202 if num.rank(data) == 0:
203 data.shape = (1,1)
204 elif num.rank(data) == 1:
205 data.shape = (1,) + data.shape
206 if num.rank(kernel) == 0:
207 kernel.shape = (1,1)
208 elif num.rank(kernel) == 1:
209 kernel.shape = (1,) + kernel.shape
210 if (kernel.shape[0] > data.shape[0] and
211 kernel.shape[1] > data.shape[1]):
212 kernel, data = data, kernel
213 elif (kernel.shape[0] <= data.shape[0] and
214 kernel.shape[1] <= data.shape[1]):
215 pass
216 return data, kernel
217
218 -def correlate2d(data, kernel, output=None, mode="nearest", cval=0.0, fft=0):
219 """correlate2d does 2d correlation of 'data' with 'kernel', storing
220 the result in 'output'.
221
222 supported 'mode's include:
223 'nearest' elements beyond boundary come from nearest edge pixel.
224 'wrap' elements beyond boundary come from the opposite array edge.
225 'reflect' elements beyond boundary come from reflection on same array edge.
226 'constant' elements beyond boundary are set to 'cval'
227
228 If fft is True, the correlation is performed using the FFT, else the
229 correlation is performed using the naive approach.
230
231 >>> a = num.arange(20*20)
232 >>> a = a.reshape((20,20))
233 >>> b = num.ones((5,5), dtype=num.float64)
234 >>> rn = correlate2d(a, b, fft=0)
235 >>> rf = correlate2d(a, b, fft=1)
236 >>> num.alltrue(num.ravel(rn-rf<1e-10))
237 True
238 """
239 data, kernel = _fix_data_kernel(data, kernel)
240 if fft:
241 return _correlate2d_fft(data, kernel, output, mode, cval)
242 else:
243 a = _correlate2d_naive(data, kernel, output, mode, cval)
244
245 return a
246
247 -def convolve2d(data, kernel, output=None, mode="nearest", cval=0.0, fft=0):
248 """convolve2d does 2d convolution of 'data' with 'kernel', storing
249 the result in 'output'.
250
251 supported 'mode's include:
252 'nearest' elements beyond boundary come from nearest edge pixel.
253 'wrap' elements beyond boundary come from the opposite array edge.
254 'reflect' elements beyond boundary come from reflection on same array edge.
255 'constant' elements beyond boundary are set to 'cval'
256
257 >>> a = num.arange(20*20)
258 >>> a = a.reshape((20,20))
259 >>> b = num.ones((5,5), dtype=num.float64)
260 >>> rn = convolve2d(a, b, fft=0)
261 >>> rf = convolve2d(a, b, fft=1)
262 >>> num.alltrue(num.ravel(rn-rf<1e-10))
263 True
264 """
265 data, kernel = _fix_data_kernel(data, kernel)
266 kernel = kernel[::-1,::-1]
267 if fft:
268 return _correlate2d_fft(data, kernel, output, mode, cval)
269 else:
270 return _correlate2d_naive(data, kernel, output, mode, cval)
271
272 -def _boxcar(data, output, boxshape, mode, cval):
273 if len(boxshape) == 1:
274 _correlate.Boxcar2d(data[num.newaxis,...], 1, boxshape[0],
275 output[num.newaxis,...], mode, cval)
276 elif len(boxshape) == 2:
277 _correlate.Boxcar2d(data, boxshape[0], boxshape[1], output, mode, cval)
278 else:
279 raise ValueError("boxshape must be a 1D or 2D shape.")
280
281 -def boxcar(data, boxshape, output=None, mode="nearest", cval=0.0):
282 """boxcar computes a 1D or 2D boxcar filter on every 1D or 2D subarray of data.
283
284 'boxshape' is a tuple of integers specifying the dimensions of the filter: e.g. (3,3)
285
286 if 'output' is specified, it should be the same shape as 'data' and
287 None will be returned.
288
289 supported 'mode's include:
290 'nearest' elements beyond boundary come from nearest edge pixel.
291 'wrap' elements beyond boundary come from the opposite array edge.
292 'reflect' elements beyond boundary come from reflection on same array edge.
293 'constant' elements beyond boundary are set to 'cval'
294
295 >>> boxcar(num.array([10, 0, 0, 0, 0, 0, 1000]), (3,), mode="nearest").astype(num.longlong)
296 array([ 6, 3, 0, 0, 0, 333, 666], dtype=int64)
297 >>> boxcar(num.array([10, 0, 0, 0, 0, 0, 1000]), (3,), mode="wrap").astype(num.longlong)
298 array([336, 3, 0, 0, 0, 333, 336], dtype=int64)
299 >>> boxcar(num.array([10, 0, 0, 0, 0, 0, 1000]), (3,), mode="reflect").astype(num.longlong)
300 array([ 6, 3, 0, 0, 0, 333, 666], dtype=int64)
301 >>> boxcar(num.array([10, 0, 0, 0, 0, 0, 1000]), (3,), mode="constant").astype(num.longlong)
302 array([ 3, 3, 0, 0, 0, 333, 333], dtype=int64)
303 >>> a = num.zeros((10,10),dtype=num.int64)
304 >>> a[0,0] = 100
305 >>> a[5,5] = 1000
306 >>> a[9,9] = 10000
307 >>> boxcar(a, (3,3)).astype(num.longlong)
308 array([[ 44, 22, 0, 0, 0, 0, 0, 0, 0, 0],
309 [ 22, 11, 0, 0, 0, 0, 0, 0, 0, 0],
310 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
311 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
312 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
313 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
314 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
315 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
316 [ 0, 0, 0, 0, 0, 0, 0, 0, 1111, 2222],
317 [ 0, 0, 0, 0, 0, 0, 0, 0, 2222, 4444]], dtype=int64)
318 >>> boxcar(a, (3,3), mode="wrap").astype(num.longlong)
319 array([[1122, 11, 0, 0, 0, 0, 0, 0, 1111, 1122],
320 [ 11, 11, 0, 0, 0, 0, 0, 0, 0, 11],
321 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
322 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
323 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
324 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
325 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
326 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
327 [1111, 0, 0, 0, 0, 0, 0, 0, 1111, 1111],
328 [1122, 11, 0, 0, 0, 0, 0, 0, 1111, 1122]], dtype=int64)
329 >>> boxcar(a, (3,3), mode="reflect").astype(num.longlong)
330 array([[ 44, 22, 0, 0, 0, 0, 0, 0, 0, 0],
331 [ 22, 11, 0, 0, 0, 0, 0, 0, 0, 0],
332 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
333 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
334 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
335 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
336 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
337 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
338 [ 0, 0, 0, 0, 0, 0, 0, 0, 1111, 2222],
339 [ 0, 0, 0, 0, 0, 0, 0, 0, 2222, 4444]], dtype=int64)
340 >>> boxcar(a, (3,3), mode="constant").astype(num.longlong)
341 array([[ 11, 11, 0, 0, 0, 0, 0, 0, 0, 0],
342 [ 11, 11, 0, 0, 0, 0, 0, 0, 0, 0],
343 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
344 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
345 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
346 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
347 [ 0, 0, 0, 0, 111, 111, 111, 0, 0, 0],
348 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
349 [ 0, 0, 0, 0, 0, 0, 0, 0, 1111, 1111],
350 [ 0, 0, 0, 0, 0, 0, 0, 0, 1111, 1111]], dtype=int64)
351
352 >>> a = num.zeros((10,10))
353 >>> a[3:6,3:6] = 111
354 >>> boxcar(a, (3,3)).astype(num.longlong)
355 array([[ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
356 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
357 [ 0, 0, 12, 24, 37, 24, 12, 0, 0, 0],
358 [ 0, 0, 24, 49, 74, 49, 24, 0, 0, 0],
359 [ 0, 0, 37, 74, 111, 74, 37, 0, 0, 0],
360 [ 0, 0, 24, 49, 74, 49, 24, 0, 0, 0],
361 [ 0, 0, 12, 24, 37, 24, 12, 0, 0, 0],
362 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
363 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
364 [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], dtype=int64)
365 """
366 mode = pix_modes[ mode ]
367 if output is None:
368 woutput = data.astype(num.float64)
369 else:
370 woutput = output
371 _fbroadcast(_boxcar, len(boxshape), data.shape,
372 (data, woutput), (boxshape, mode, cval))
373 if output is None:
374 return woutput
375
377 """_fbroadcast(f, N, args, shape, params=()) calls 'f' for each of the
378 'N'-dimensional inner subnumarray of 'args'. Each subarray has
379 .shape == 'shape'[-N:]. There are a total of product(shape[:-N],axis=0)
380 calls to 'f'.
381 """
382 if len(shape) == N:
383 apply(f, tuple(args)+params)
384 else:
385 for i in range(shape[0]):
386 _fbroadcast(f, N, shape[1:], [x[i] for x in args], params)
387