Эффективно получать пары положительных значений с numpy

Question

У меня есть функция python, которая принимает два списка, ищет пары на двух входах, где оба имеют положительные значения в одном и том же индексе, и создает два выходных списка, добавляя к каждому из них эти два положительных значения. У меня есть рабочая функция:

def get_pairs_in_first_quadrant(x_in, y_in): 
    """If both x_in[i] and y_in[i] are > 0 then both will appended to the output list. If either are negative 
    then the pair of them will be absent from the output list. 
    :param x_in: A list of positive or negative floats 
    :param y_in: A list of positive or negative floats 
    :return: A list of positive floats <= in length to the inputs. 
    """ 
    x_filtered, y_filtered = [], [] 
    for x, y in zip(x_in, y_in): 
     if x > 0 and y > 0: 
      x_filtered.append(x) 
      y_filtered.append(y) 
    return x_filtered, y_filtered 

Как сделать это быстрее, используя numpy?

Answer 1

Вы можете сделать это, просто найти индексы, где они оба положительны:

import numpy as np 

a = np.random.random(10) - .5 
b = np.random.random(10) - .5 

def get_pairs_in_first_quadrant(x_in, y_in): 
    i = np.nonzero((x_in>0) & (y_in>0)) # main line of interest 
    return x_in[i], y_in[i] 

print a # [-0.18012451 -0.40924713 -0.3788772 0.3186816 0.14811581 -0.04021951 -0.21278312 -0.36762629 -0.45369899 -0.46374929] 
print b # [ 0.33005969 -0.03167875 0.11387641 0.22101336 0.38412264 -0.3880842 0.08679424 0.3126209 -0.08760505 -0.40921421] 
print get_pairs_in_first_quadrant(a, b) # (array([ 0.3186816 , 0.14811581]), array([ 0.22101336, 0.38412264])) 

---

Я был заинтересован в предложении Хайма просто используя логическое индексирование без вызова nonzero, так что я провел несколько тестов синхронизации. Результаты несколько интересны, поскольку соотношение преимуществ немонотонное с количеством положительных совпадений, но в основном, по крайней мере, для скорости, на самом деле не имеет значения, что используется (хотя nonzero обычно немного быстрее и может быть около в два раза быстрее):

threshold = .6 
a = np.random.random(10000) - threshold 
b = np.random.random(10000) - threshold 

def f1(x_in, y_in): 
    i = np.nonzero((x_in>0) & (y_in>0)) # main line of interest 
    return x_in[i], y_in[i] 

def f2(x_in, y_in): 
    i = (x_in>0) & (y_in>0) # main line of interest 
    return x_in[i], y_in[i] 

print threshold, len(f1(a,b)[0]), len(f2(a,b)[0]) 
print timeit("f1(a, b)", "from __main__ import a, b, f1, f2", number = 1000) 
print timeit("f2(a, b)", "from __main__ import a, b, f1, f2", number = 1000) 

Который дает для различных пороговых значений:

0.05 9086 9086 
0.0815141201019 
0.104746818542 

0.5 2535 2535 
0.0715141296387 
0.153401851654 

0.95 21 21 
0.027126789093 
0.0324990749359