10000 ENH: add maxlag to np.correlate & convolve (gh-5954, gh-4940, gh-5978) · numpy/numpy@e7228b0 · GitHub
[go: up one dir, main page]
Skip to content

Commit e7228b0

Browse files
HoniSandersHoniSanders
committed
ENH: add maxlag to np.correlate & convolve (gh-5954, gh-4940, gh-5978)
1 parent c10f7d0 commit e7228b0

File tree

2 files changed

+217
-23
lines changed

2 files changed

+217
-23
lines changed

‎numpy/core/numeric.py

Lines changed: 116 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
from . import umath
2525
from .umath import (multiply, invert, sin, UFUNC_BUFSIZE_DEFAULT,
2626
ERR_IGNORE, ERR_WARN, ERR_RAISE, ERR_CALL, ERR_PRINT,
27-
ERR_LOG, ERR_DEFAULT, PINF, NAN)
27+
ERR_LOG, ERR_DEFAULT, PINF, NAN, ceil)
2828
from . import numerictypes
2929
from .numerictypes import longlong, intc, int_, float_, complex_, bool_
3030
from ._internal import TooHardError, AxisError
@@ -907,7 +907,35 @@ def _mode_from_name(mode):
907907
return mode
908908

909909

910-
def correlate(a, v, mode='valid'):
910+
def _lags_from_mode(alen, vlen, mode):
911+
if type(mode) is int: # maxlag
912+
lags = (-mode+1, mode, 1)
913+
mode = 3
914+
elif type(mode) is tuple: # minlag and maxlag
915+
if len(mode) > 2:
916+
lags = (int(mode[0]), int(mode[1]), int(mode[2]))
917+
else:
918+
lags = (int(mode[0]), int(mode[1]), 1)
919+
mode = 3
920+
elif isinstance(mode, basestring):
921+
mode = _mode_from_name_dict[mode.lower()[0]]
922+
if alen < vlen:
923+
alen, vlen = vlen, alen
924+
inverted = 1
925+
else:
926+
inverted = 0
927+
if mode is 0:
928+
lags = (0, alen-vlen+1, 1)
929+
elif mode is 1:
930+
lags = (-int(vlen/2), alen - int(vlen/2), 1)
931+
elif mode is 2:
932+
lags = (-vlen+1, alen, 1)
933+
if inverted:
934+
lags = (-int(ceil((lags[1]-lags[0])/float(lags[2])))*lags[2]-lags[0]+lags[2], -lags[0]+lags[2], lags[2])
935+
return mode, lags
936+
937+
938+
def correlate(a, v, mode='valid', returns_lags=False):
911939
"""
912940
Cross-correlation of two 1-dimensional sequences.
913941
@@ -923,9 +951,15 @@ def correlate(a, v, mode='valid'):
923951
----------
924952
a, v : array_like
925953
Input sequences.
926-
mode : {'valid', 'same', 'full'}, optional
954+
mode : int, int tuple, or {'valid', 'same', 'full'}, optional
927955
Refer to the `convolve` docstring. Note that the default
928-
is 'valid', unlike `convolve`, which uses 'full'.
956+
is `valid`, unlike `convolve`, which uses `full`.
957+
returns_lags : bool, optional
958+
If True, the function returns a lagvector array in addition to the
959+
cross-correlation result. The lagvector contains the indices of
960+
the lags for which the cross-correlation was calculated. It is
961+
the same length as the return array, and corresponds one-to-one.
962+
False is default.
929963
old_behavior : bool
930964
`old_behavior` was removed in NumPy 1.10. If you need the old
931965
behavior, use `multiarray.correlate`.
@@ -934,6 +968,9 @@ def correlate(a, v, mode='valid'):
934968
-------
935969
out : ndarray
936970
Discrete cross-correlation of `a` and `v`.
971+
lagvector : ndarray, optional
972+
The indices of the lags for which the cross-correlation was calculated.
973+
It is the same length as out, and corresponds one-to-one.
937974
938975
See Also
939976
--------
@@ -953,10 +990,14 @@ def correlate(a, v, mode='valid'):
953990
--------
954991
>>> np.correlate([1, 2, 3], [0, 1, 0.5])
955992
array([ 3.5])
956-
>>> np.correlate([1, 2, 3], [0, 1, 0.5], "same")
993+
>>> np.correlate([1, 2, 3], [0, 1, 0.5], mode="same")
994+
array([ 2. , 3.5, 3. ])
995+
>>> np.correlate([1, 2, 3], [0, 1, 0.5], mode="full", returns_lags=True)
996+
(array([ 0.5, 2. , 3.5, 3. , 0. ]), array([-2, -1, 0, 1, 2]))
997+
>>> np.correlate([1, 2, 3], [0, 1, 0.5], mode=2)
957998
array([ 2. , 3.5, 3. ])
958-
>>> np.correlate([1, 2, 3], [0, 1, 0.5], "full")
959-
array([ 0.5, 2. , 3.5, 3. , 0. ])
999+
>>> np.correlate([1, 2, 3], [0, 1, 0.5], mode=(-1,2,2), returns_lags=True)
1000+
(array([ 2., 3.]), array([-1, 1]))
9601001
9611002
Using complex sequences:
9621003
@@ -971,11 +1012,15 @@ def correlate(a, v, mode='valid'):
9711012
array([ 0.0+0.j , 3.0+1.j , 1.5+1.5j, 1.0+0.j , 0.5+0.5j])
9721013
9731014
"""
974-
mode = _mode_from_name(mode)
975-
return multiarray.correlate2(a, v, mode)
1015+
mode, lags = _lags_from_mode(len(a), len(v), mode)
1016+
if returns_lags:
1017+
return multiarray.correlate2(a, v, 3, lags[0], lags[1], lags[2]), \
1018+
arange(lags[0], lags[1], lags[2])
1019+
else:
1020+
return multiarray.correlate2(a, v, 3, lags[0], lags[1], lags[2])
9761021

9771022

978-
def convolve(a, v, mode='full'):
1023+
def convolve(a, v, mode='full', returns_lags=False):
9791024
"""
9801025
Returns the discrete, linear convolution of two one-dimensional sequences.
9811026
@@ -993,27 +1038,53 @@ def convolve(a, v, mode='full'):
9931038
First one-dimensional input array.
9941039
v : (M,) array_like
9951040
Second one-dimensional input array.
996-
mode : {'full', 'valid', 'same'}, optional
1041+
mode : int, int tuple, or {'full', 'valid', 'same'}, optional
1042+
int (maxlag):
1043+
This calculates the convolution for all lags starting at
1044+
(-maxlag + 1) and ending at (maxlag - 1), with steps of size 1.
1045+
See the optional lagvec argument to get an array containing
1046+
lags corresponding to the convolution values in the return array.
1047+
1048+
tuple (minlag, maxlag) or (minlag, maxlag, lagstep):
1049+
This calculates the convolution for all lags starting at
1050+
minlag and ending at (maxlag - 1), with steps of size lagstep.
1051+
The lags for which the convolution will be calculated correspond
1052+
with the values in the vector formed by numpy.arange() with the
1053+
same tuple argument.
1054+
9971055
'full':
9981056
By default, mode is 'full'. This returns the convolution
9991057
at each point of overlap, with an output shape of (N+M-1,). At
10001058
the end-points of the convolution, the signals do not overlap
1001-
completely, and boundary effects may be seen.
1059+
completely, and boundary effects may be seen. This corresponds
1060+
with a lag tuple of (-M+1, N, 1) for N>M or (-N+1, M, 1)
1061+
for M>N.
10021062
10031063
'same':
1004-
Mode 'same' returns output of length ``max(M, N)``. Boundary
1005-
effects are still visible.
1064+
Mode `same` returns output of length ``max(M, N)``. Boundary
1065+
effects are still visible. This corresponds with a lag tuple of
1066+
(-M/2, N-M/2, 1) for N>M or (-M+N/2+1, N/2+1, 1) for M>N.
10061067
10071068
'valid':
10081069
Mode 'valid' returns output of length
10091070
``max(M, N) - min(M, N) + 1``. The convolution product is only given
10101071
for points where the signals overlap completely. Values outside
1011-
the signal boundary have no effect.
1072+
the signal boundary have no effect. This corresponds with a lag tuple
1073+
of (0, N-M+1, 1) for N>M or (-M+N, 1, 1) for M>N.
1074+
returns_lags : bool, optional
1075+
If True, the function returns a lagvector array in addition to the
1076+
convolution result. The lagvector contains the indices of
1077+
the lags for which the convolution was calculated. It is
1078+
the same length as the return array, and corresponds one-to-one.
1079+
False is default.
10121080
10131081
Returns
10141082
-------
10151083
out : ndarray
10161084
Discrete, linear convolution of `a` and `v`.
1085+
lagvector : ndarray, optional
1086+
The indices of the lags for which the convolution was calculated.
1087+
It is the same length as out, and corresponds one-to-one.
10171088
10181089
See Also
10191090
--------
@@ -1052,14 +1123,34 @@ def convolve(a, v, mode='full'):
10521123
Contains boundary effects, where zeros are taken
10531124
into account:
10541125
1055-
>>> np.convolve([1,2,3],[0,1,0.5], 'same')
1126+
>>> np.convolve([1,2,3],[0,1,0.5], mode='same')
10561127
array([ 1. , 2.5, 4. ])
10571128
10581129
The two arrays are of the same length, so there
1059-
is only one position where they completely overlap:
1130+
is only one position where they completely overlap,
1131+
corresponding to a lag of 0. lagvector=True causes
1132+
the function to return the lagvector corresponding
1133+
to the convolution in addition to the convolution
1134+
itself:
1135+
1136+
>>> np.convolve([1,2,3],[0,1,0.5], mode='valid', returns_lags=True)
1137+
(array([ 2.5]), array([0]))
1138+
1139+
Find the convolution for lags ranging from -1 to 1
1140+
(0 is the lag for which the left sides of the arrays
1141+
are aligned, -1 has the second vector to the left of
1142+
the first, and +1 has the second vector to the right
1143+
of the first):
10601144
1061-
>>> np.convolve([1,2,3],[0,1,0.5], 'valid')
1062-
array([ 2.5])
1145+
>>> np.convolve([1,2,3],[0,1,0.5], mode=2, returns_lags=True)
1146+
(array([ 1. , 2.5, 4. ]), array([-1, 0, 1]))
1147+
1148+
Find the convolution for lags ranging from -2 to 4
1149+
with steps of length 2 (the maxlag member of the
1150+
lag range tuple is non-inclusive, similar to np.arange()):
1151+
1152+
>>> np.convolve([1,2,3,4,5],[0,1,0.5], mode=(-2,6,2), returns_lags=True)
1153+
(array([ 0. , 2.5, 5.5, 2.5]), array([-2, 0, 2, 4]))
10631154
10641155
"""
10651156
a, v = array(a, copy=False, ndmin=1), array(v, copy=False, ndmin=1)
@@ -1069,8 +1160,12 @@ def convolve(a, v, mode='full'):
10691160
raise ValueError('a cannot be empty')
10701161
if len(v) == 0:
10711162
raise ValueError('v cannot be empty')
1072-
mode = _mode_from_name(mode)
1073-
return multiarray.correlate(a, v[::-1], mode)
1163+
mode, lags = _lags_from_mode(len(a), len(v), mode)
1164+
if returns_lags:
1165+
return multiarray.correlate2(a, v[::-1], 3, lags[0], lags[1], lags[2]), \
1166+
arange(lags[0], lags[1], lags[2])
1167+
else:
1168+
return multiarray.correlate2(a, v[::-1], 3, lags[0], lags[1], lags[2])
10741169

10751170

10761171
def outer(a, b, out=None):

‎numpy/core/tests/test_numeric.py

Lines changed: 101 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2177,28 +2177,127 @@ def _setup(self, dt):
21772177
self.xs = np.arange(1, 20)[::3]
21782178
self.y = np.array([-1, -2, -3], dtype=dt)
21792179
self.z1 = np.array([ -3., -8., -14., -20., -26., -14., -5.], dtype=dt)
2180+
self.z1s = np.array([-8., -14., -20., -26., -14.], dtype=dt)
2181+
self.z1v = np.array([-14., -20., -26.], dtype=dt)
21802182
self.z1_4 = np.array([-2., -5., -8., -11., -14., -5.], dtype=dt)
21812183
self.z1r = np.array([-15., -22., -22., -16., -10., -4., -1.], dtype=dt)
21822184
self.z2 = np.array([-5., -14., -26., -20., -14., -8., -3.], dtype=dt)
2185+
self.z2s = np.array([-14, -26., -20., -14., -8.], dtype=dt)
2186+
self.z2v = np.array([-26., -20., -14.], dtype=dt)
21832187
self.z2r = np.array([-1., -4., -10., -16., -22., -22., -15.], dtype=dt)
21842188
self.zs = np.array([-3., -14., -30., -48., -66., -84.,
21852189
-102., -54., -19.], dtype=dt)
21862190

21872191
def test_float(self):
21882192
self._setup(np.float)
2189-
z = np.correlate(self.x, self.y, 'full')
2193+
z, lagvec = np.correlate(self.x, self.y, 'full', returns_lags=True)
21902194
assert_array_almost_equal(z, self.z1)
2195+
assert_equal(lagvec.size, z.size)
2196+
assert_array_equal(lagvec, np.arange(-self.y.size + 1, self.x.size))
2197+
z, lagvec = np.correlate(self.x, self.y, 'same', returns_lags=True)
2198+
assert_array_almost_equal(z, self.z1s)
2199+
assert_array_almost_equal(z, self.z1[lagvec+self.y.size-1])
2200+
assert_equal(lagvec.size, z.size)
2201+
if self.x.size > self.y.size:
2202+
assert_array_equal(lagvec, np.arange(-int(self.y.size/2), self.x.size - int(self.y.size/2)))
2203+
else:
2204+
assert_array_equal(lagvec, np.arange(-self.y.size + int(self.x.size/2) + 1, int(self.x.size/2) + 1))
2205+
z, lagvec = np.correlate(self.x, self.y, 'valid', returns_lags=True)
2206+
assert_array_almost_equal(z, self.z1v)
2207+
assert_array_almost_equal(z, self.z1[lagvec+self.y.size-1])
2208+
assert_equal(lagvec.size, z.size)
2209+
if self.x.size > self.y.size:
2210+
assert_array_equal(lagvec, np.arange(0, self.x.size - self.y.size + 1))
2211+
else:
2212+
assert_array_equal(lagvec, np.arange(self.x.size - self.y.size, 1))
2213+
21912214
z = np.correlate(self.x, self.y[:-1], 'full')
21922215
assert_array_almost_equal(z, self.z1_4)
2193-
z = np.correlate(self.y, self.x, 'full')
2216+
z, lagvec = np.correlate(self.y, self.x, 'full', returns_lags=True)
21942217
assert_array_almost_equal(z, self.z2)
2218+
assert_equal(lagvec.size, z.size)
2219+
assert_array_equal(lagvec, np.arange(-self.x.size + 1, self.y.size))
2220+
z, lagvec = np.correlate(self.y, self.x, 'same', returns_lags=True)
2221+
assert_array_almost_equal(z, self.z2s)
2222+
assert_array_almost_equal(z, self.z2[lagvec+self.x.size-1])
2223+
assert_equal(lagvec.size, z.size)
2224+
if self.y.size > self.x.size:
2225+
assert_array_equal(lagvec, np.arange(-int(self.x.size/2), self.y.size - int(self.x.size/2)))
2226+
else:
2227+
assert_array_equal(lagvec, np.arange(-self.x.size + int(self.y.size/2) + 1, int(self.y.size/2) + 1))
2228+
z, lagvec = np.correlate(self.y, self.x, 'valid', returns_lags=True)
2229+
assert_array_almost_equal(z, self.z2v)
2230+
assert_array_almost_equal(z, self.z2[lagvec+self.x.size-1])
2231+
assert_equal(lagvec.size, z.size)
2232+
if self.y.size > self.x.size:
2233+
assert_array_equal(lagvec, np.arange(0, self.y.size - self.x.size + 1))
2234+
else:
2235+
assert_array_equal(lagvec, np.arange(self.y.size - self.x.size, 1))
21952236
z = np.correlate(self.x[::-1], self.y, 'full')
21962237
assert_array_almost_equal(z, self.z1r)
21972238
z = np.correlate(self.y, self.x[::-1], 'full')
21982239
assert_array_almost_equal(z, self.z2r)
21992240
z = np.correlate(self.xs, self.y, 'full')
22002241
assert_array_almost_equal(z, self.zs)
22012242

2243+
def test_lags(self):
2244+
self._setup(np.float)
2245+
longlen = max(self.x.size, self.y.size)
2246+
maxlag = 3 + longlen
2247+
minlag = -2
2248+
lagstep = 2
2249+
tmp = np.correlate(self.x, self.y, 'full')
2250+
z_full = np.zeros(tmp.size + 100)
2251+
z_full[:tmp.size] = tmp
2252+
z, lagvec = np.correlate(self.x, self.y, maxlag, returns_lags=True)
2253+
assert_array_equal(z[0:3], np.zeros(3))
2254+
assert_array_equal(z[-3:], np.zeros(3))
2255+
assert_equal(lagvec.size, z.size)
2256+
assert_array_equal(lagvec, np.arange(-maxlag+1, maxlag))
2257+
z, lagvec = np.correlate(self.x, self.y, (minlag, maxlag), returns_lags=True)
2258+
assert_array_almost_equal(z, z_full[lagvec+self.y.size-1])
2259+
assert_equal(lagvec.size, z.size)
2260+
assert_array_equal(lagvec, np.arange(minlag, maxlag))
2261+
z, lagvec = np.correlate(self.x, self.y, (minlag, maxlag, lagstep), returns_lags=True)
2262+
assert_array_almost_equal(z, z_full[lagvec+self.y.size-1])
2263+
assert_equal(lagvec.size, z.size)
2264+
assert_array_equal(lagvec, np.arange(minlag, maxlag, lagstep))
2265+
lagstep = 3
2266+
z, lagvec = np.correlate(self.x, self.y, (minlag, maxlag, lagstep), returns_lags=True)
2267+
assert_array_almost_equal(z, z_full[lagvec+self.y.size-1])
2268+
assert_equal(lagvec.size, z.size)
2269+
minlag = 10
2270+
maxlag = -2
2271+
lagstep = -2
2272+
z, lagvec = np.correlate(self.x, self.y, (minlag, maxlag, lagstep), returns_lags=True)
2273+
assert_array_almost_equal(z, z_full[lagvec+self.y.size-1])
2274+
assert_equal(lagvec.size, z.size)
2275+
assert_array_equal(lagvec, np.arange(int(minlag), int(maxlag), lagstep))
2276+
maxlag = 10.2
2277+
minlag = -2.2
2278+
z, lagvec = np.correlate(self.x, self.y, (minlag, maxlag), returns_lags=True)
2279+
assert_array_almost_equal(z, z_full[lagvec+self.y.size-1])
2280+
assert_equal(lagvec.size, z.size)
2281+
assert_array_equal(lagvec, np.arange(int(minlag), int(maxlag)))
2282+
z, lagvec = np.correlate(self.x, self.y, (maxlag, minlag), returns_lags=True)
2283+
assert_array_almost_equal(z, z_full[lagvec+self.y.size-1])
2284+
assert_equal(lagvec.size, z.size)
2285+
assert_array_equal(lagvec, np.arange(int(maxlag), int(minlag)))
2286+
maxlag = -1
2287+
minlag = -2
2288+
z, lagvec = np.correlate(self.x, self.y, (minlag, maxlag), returns_lags=True)
2289+
assert_array_almost_equal(z, z_full[lagvec+self.y.size-1])
2290+
assert_equal(lagvec.size, z.size)
2291+
assert_array_equal(lagvec, np.arange(int(minlag), int(maxlag)))
2292+
maxlag = 2
2293+
minlag = 4
2294+
z, lagvec = np.correlate(self.x, self.y, (minlag, maxlag), returns_lags=True)
2295+
assert_array_almost_equal(z, z_full[lagvec+self.y.size-1])
2296+
assert_equal(lagvec.size, z.size)
2297+
assert_array_equal(lagvec, np.arange(int(minlag), int(maxlag)))
2298+
# check non-integer lagsteps?
2299+
# i want to be able to switch x and y and test
2300+
22022301
def test_object(self):
22032302
self._setup(Decimal)
22042303
z = np.correlate(self.x, self.y, 'full')

0 commit comments

Comments
 (0)
0