← Back to branch summary

~ubuntu-branches/ubuntu/saucy/python-scipy/saucy

« back to all changes in this revision

Viewing changes to scipy/stsci/convolve/lib/Convolve.py

  • Committer: Bazaar Package Importer
  • Author(s): Ondrej Certik
  • Date: 2008-06-16 22:58:01 UTC
  • mfrom: (2.1.24 intrepid)
  • Revision ID: james.westby@ubuntu.com-20080616225801-irdhrpcwiocfbcmt
Tags: 0.6.0-12
http://bugs.debian.org/489149
* The description updated to match the current SciPy (Closes: #489149).
* Standards-Version bumped to 3.8.0 (no action needed)
* Build-Depends: netcdf-dev changed to libnetcdf-dev

Show diffs side-by-side

added added

removed removed

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