Numpy数组条件匹配

Question

我需要对比两个非常大的Numpy数组，一个有20000行，另一个大约有100000行。我正在尝试写一个脚本来高效地完成这个任务。简单地逐行比较数组速度非常慢，有人能建议更好的方法吗？我想做的是：数组 datesSecondDict 和数组 pwfs2Dates 里都有日期时间值，我需要从数组 pwfs2Dates（较小的数组）中取出每个日期时间值，看看在数组 datesSecondDict 中是否有类似的日期时间值（加减5分钟都算）。如果找到了一个（或者多个），我就用数组 valsSecondDict 中对应的数值填充一个新数组（这个新数组和 pwfs2Dates 的大小一样）。这是 @unutbu 和 @joaquin 提供的一个解决方案，对我有用（谢谢你们！）：

import time
import datetime as dt
import numpy as np

def combineArs(dict1, dict2):
   """Combine data from 2 dictionaries into a list.
   dict1 contains primary data (e.g. seeing parameter).
   The function compares each timestamp in dict1 to dict2
   to see if there is a matching timestamp record(s)
   in dict2 (plus/minus 5 minutes).
   ==If yes: a list called data gets appended with the
   corresponding parameter value from dict2.
   (Note that if there are more than 1 record matching,
   the first occuring value gets appended to the list).
   ==If no: a list called data gets appended with 0."""
   # Specify the keys to use    
   pwfs2Key = 'pwfs2:dc:seeing'
   dimmKey = 'ws:seeFwhm'

   # Create an iterator for primary dict 
   datesPrimDictIter = iter(dict1[pwfs2Key]['datetimes'])

   # Take the first timestamp value in primary dict
   nextDatePrimDict = next(datesPrimDictIter)

   # Split the second dictionary into lists
   datesSecondDict = dict2[dimmKey]['datetime']
   valsSecondDict  = dict2[dimmKey]['values']

   # Define time window
   fiveMins = dt.timedelta(minutes = 5)
   data = []
   #st = time.time()
   for i, nextDateSecondDict in enumerate(datesSecondDict):
       try:
           while nextDatePrimDict < nextDateSecondDict - fiveMins:
               # If there is no match: append zero and move on
               data.append(0)
               nextDatePrimDict = next(datesPrimDictIter)
           while nextDatePrimDict < nextDateSecondDict + fiveMins:
               # If there is a match: append the value of second dict
               data.append(valsSecondDict[i])
               nextDatePrimDict = next(datesPrimDictIter)
       except StopIteration:
           break
   data = np.array(data)   
   #st = time.time() - st    
   return data

谢谢，
Aina。

Answer 1

我觉得你可以用少一个循环来实现这个功能：

import datetime
import numpy

# Test data

# Create an array of dates spaced at 1 minute intervals
m = range(1, 21)
n = datetime.datetime.now()
a = numpy.array([n + datetime.timedelta(minutes=i) for i in m])

# A smaller array with three of those dates
m = [5, 10, 15]
b = numpy.array([n + datetime.timedelta(minutes=i) for i in m])

# End of test data

def date_range(date_array, single_date, delta):
    plus = single_date + datetime.timedelta(minutes=delta)
    minus = single_date - datetime.timedelta(minutes=delta)
    return date_array[(date_array < plus) * (date_array > minus)]

dates = []
for i in b:
    dates.append(date_range(a, i, 5))

all_matches = numpy.unique(numpy.array(dates).flatten())

肯定还有更好的方法来收集和合并匹配的结果，不过你明白我的意思了……你也可以使用 numpy.argwhere((a < plus) * (a > minus)) 来返回索引，而不是日期，然后用这个索引来获取整行数据，并把它放到你的新数组里。

Answer 2

基于joaquin的想法：

import datetime as dt
import itertools

def combineArs(dict1, dict2, delta = dt.timedelta(minutes = 5)):
    marks = dict1['datetime']
    values = dict1['values']
    pdates = iter(dict2['datetime'])

    data = []
    datei = next(pdates)
    for datej, val in itertools.izip(marks, values):
        try:
            while datei < datej - delta:
                data.append(0)
                datei = next(pdates)
            while datei < datej + delta:
                data.append(val)
                datei = next(pdates)
        except StopIteration:
            break
    return data

dict1 = { 'ws:seeFwhm':
          {'datetime': [dt.datetime(2011, 12, 19, 12, 0, 0),
                        dt.datetime(2011, 12, 19, 12, 1, 0),
                        dt.datetime(2011, 12, 19, 12, 20, 0),
                        dt.datetime(2011, 12, 19, 12, 22, 0),
                        dt.datetime(2011, 12, 19, 12, 40, 0), ],
           'values': [1, 2, 3, 4, 5] } }
dict2 = { 'pwfs2:dc:seeing':
          {'datetime': [dt.datetime(2011, 12, 19, 12, 9),
                         dt.datetime(2011, 12, 19, 12, 19),
                         dt.datetime(2011, 12, 19, 12, 29),
                         dt.datetime(2011, 12, 19, 12, 39),
                        ], } }

if __name__ == '__main__':
    dimmKey = 'ws:seeFwhm'
    pwfs2Key = 'pwfs2:dc:seeing'    
    print(combineArs(dict1[dimmKey], dict2[pwfs2Key]))

会产生

[0, 3, 0, 5]

Answer 3

数组中的日期是排好序的吗？

• 如果是的话，你可以通过在内层循环中，一旦发现日期比外层循环给出的日期大，就直接跳出比较，这样可以加快比较的速度。这样你就只需要一次比较，而不是要循环 dimVals 项目 len(pwfs2Vals) 次。
• 如果不是，或许你应该把当前的 pwfs2Dates 数组转换成，比如说，成对的数组 [(date, array_index),...]，然后你可以先按日期对所有数组进行排序，这样就能进行上面提到的一次性比较，同时也能得到设置 data[i] 所需的原始索引。

例如，如果这些数组已经排好序（我这里用的是列表，不确定你是否需要数组）：

pdates = iter(enumerate(pwfs2Dates))
i, datei = pdates.next() 

for datej, valuej in zip(dimmDates, dimvals):
    while datei < datej - fiveMinutes:
        i, datei = pdates.next()
    while datei < datej + fiveMinutes:
        data[i] = valuej
        i, datei = pdates.next()

否则，如果它们没有排序，而你创建了像这样的已排序、带索引的列表：

pwfs2Dates = sorted([(date, idx) for idx, date in enumerate(pwfs2Dates)])
dimmDates = sorted([(date, idx) for idx, date in enumerate(dimmDates)])

代码会是：
（编辑：现在使用迭代器，而不是在每一步都从头循环 pwfs2Dates）：

pdates = iter(pwfs2Dates)
datei, i = pdates.next()

for datej, j in dimmDates:
    while datei < datej - fiveMinutes:
        datei, i = pdates.next()
    while datei < datej + fiveMinutes:
        data[i] = dimVals[j]
        datei, i = pdates.next()

太好了！

..

1. 注意 dimVals：

   dimVals  = np.array(dict1[dimmKey]['values'])
   

   在你的代码中没有使用，可以去掉。

2. 注意，通过直接循环数组本身来简化代码，而不是使用 xrange。

---

编辑：来自 unutbu 的回答提到了上面代码中的一些薄弱环节。我在这里列出以便完整：

1. 使用 next：使用 next(iterator) 比 iterator.next() 更好。iterator.next() 是一个命名规则的例外，在 py3k 中已经修正，将这个方法重命名为 iterator.__next__()。
2. 用 try/except 检查迭代器的结束。迭代器中的所有项目完成后，下一次调用 next() 会产生一个 StopIteration 异常。使用 try/except 可以在这种情况下友好地跳出循环。对于 OP 问题的具体情况，这不是问题，因为两个数组大小相同，所以 for 循环和迭代器同时结束。因此不会引发异常。然而，可能会有 dict1 和 dict2 大小不相同的情况。在这种情况下，可能会引发异常。问题是：使用 try/except 更好，还是在循环之前通过将它们等同于较短的数组来准备数组更好。