Python相当于MATLAB的“ismember”函数

时间:2013-04-07 15:25:54

标签: python matlab optimization numpy

在尝试优化代码的许多尝试之后,似乎最后一个资源是尝试使用多个核运行下面的代码。我不知道如何转换/重新构建我的代码,以便使用多个内核可以更快地运行。如果我能获得指导以实现最终目标,我将不胜感激。最终目标是能够尽快为阵列A和B运行此代码,其中每个阵列包含大约700,000个元素。这是使用小数组的代码。 700k元素数组已被注释掉。

import numpy as np

def ismember(a,b):
    for i in a:
        index = np.where(b==i)[0]
        if index.size == 0:
            yield 0
        else:
            yield index


def f(A, gen_obj):
    my_array = np.arange(len(A))
    for i in my_array:
        my_array[i] = gen_obj.next()
    return my_array


#A = np.arange(700000)
#B = np.arange(700000)
A = np.array([3,4,4,3,6])
B = np.array([2,5,2,6,3])

gen_obj = ismember(A,B)

f(A, gen_obj)

print 'done'
# if we print f(A, gen_obj) the output will be: [4 0 0 4 3]
# notice that the output array needs to be kept the same size as array A.

我要做的是模仿一个名为ismember [2]的MATLAB函数(格式为:[Lia,Locb] = ismember(A,B)的函数。我只是试图得到Locb仅部分。

  

来自Matlab:Locb,包含A中每个值的B中最低的索引.A是B的成员。输出数组Locb在A不是B的成员的地方包含0

主要问题之一是我需要能够尽可能高效地执行此操作。为了测试,我有两个700k元素的数组。创建生成器并遍历生成器的值似乎不能快速完成工作。

5 个答案:

答案 0 :(得分:16)

在担心多核之前,我会使用字典消除你的ismember函数中的线性扫描:

def ismember(a, b):
    bind = {}
    for i, elt in enumerate(b):
        if elt not in bind:
            bind[elt] = i
    return [bind.get(itm, None) for itm in a]  # None can be replaced by any other "not in b" value

您的原始实现需要对A中的每个元素进行B中元素的完整扫描,使其成为O(len(A)*len(B))。上面的代码需要对B进行一次全扫描才能生成字典Bset。通过使用dict,您可以有效地为A的每个元素查找B中的每个元素,并进行操作O(len(A)+len(B))。如果这仍然太慢,那么担心在多个核心上运行上述功能。

编辑:我也略微修改了你的索引。 Matlab使用0,因为它的所有数组都从索引1开始.Python / numpy start数组为0,所以如果你的数据集看起来像这样

A = [2378, 2378, 2378, 2378]
B = [2378, 2379]

并且你没有元素返回0,那么你的结果将排除A的所有元素。上面的例程返回None没有索引而不是0.返回-1是一个选项,但是Python会将其解释为是数组中的最后一个元素。 None如果将其用作数组的索引,则会引发异常。如果您想要不同的行为,请将Bind.get(item,None)表达式中的第二个参数更改为您想要返回的值。

答案 1 :(得分:11)

sfstewman的优秀答案很可能为您解决了这个问题。

我只想添加如何在numpy中完成同样的工作。

我使用了numpy的uniquein1d函数。

B_unique_sorted, B_idx = np.unique(B, return_index=True)
B_in_A_bool = np.in1d(B_unique_sorted, A, assume_unique=True)
  • B_unique_sorted包含B已排序的唯一值。
  • B_idx为这些值保留原始B
  • 中的索引
  • B_in_A_bool是一个大小为B_unique_sorted的布尔数组 存储B_unique_sorted中的值是否在A注意:我需要在A 中查找(来自B的唯一值),因为我需要根据B_idx返回输出。 注意:我认为A已经是唯一的。

现在您可以使用B_in_A_bool获取常用值

B_unique_sorted[B_in_A_bool]

及其原始B

中的各自指数
B_idx[B_in_A_bool]

最后,我认为这比纯Python for循环要快得多,尽管我没有测试它。

答案 2 :(得分:1)

尝试使用列表理解;

In [1]: import numpy as np

In [2]: A = np.array([3,4,4,3,6])

In [3]: B = np.array([2,5,2,6,3])

In [4]: [x for x in A if not x in B]
Out[4]: [4, 4]

通常,列表推导比for循环快得多。

获得相等的长度列表;

In [19]: map(lambda x: x if x not in B else False, A)
Out[19]: [False, 4, 4, False, False]

这对于小型数据集来说非常快:

In [20]: C = np.arange(10000)

In [21]: D = np.arange(15000, 25000)

In [22]: %timeit map(lambda x: x if x not in D else False, C)
1 loops, best of 3: 756 ms per loop

对于大型数据集,您可以尝试使用multiprocessing.Pool.map()来加速操作。

答案 3 :(得分:1)

这是确切的MATLAB等价物,它返回匹配MATLAB的输出参数[Lia,Locb],除了Python 0也是一个有效的索引。所以,这个函数不会返回0。它基本上返回Locb(Locb> 0)。性能也等同于MATLAB。

def ismember(a_vec, b_vec):
    """ MATLAB equivalent ismember function """

    bool_ind = np.isin(a_vec,b_vec)
    common = a[bool_ind]
    common_unique, common_inv  = np.unique(common, return_inverse=True)     # common = common_unique[common_inv]
    b_unique, b_ind = np.unique(b_vec, return_index=True)  # b_unique = b_vec[b_ind]
    common_ind = b_ind[np.isin(b_unique, common_unique, assume_unique=True)]
    return bool_ind, common_ind[common_inv]

有点(~5x)慢但没有使用独特功能的替代实现在这里:

def ismember(a_vec, b_vec):
    ''' MATLAB equivalent ismember function. Slower than above implementation'''
    b_dict = {b_vec[i]: i for i in range(0, len(b_vec))}
    indices = [b_dict.get(x) for x in a_vec if b_dict.get(x) is not None]
    booleans = np.in1d(a_vec, b_vec)
    return booleans, np.array(indices, dtype=int)

答案 4 :(得分:1)

尝试使用ismember库。

pip install ismember

简单的例子:

# Import library
from ismember import ismember
import numpy as np

# data
A = np.array([3,4,4,3,6])
B = np.array([2,5,2,6,3])

# Lookup
Iloc,idx = ismember(A, B)
 
# Iloc is boolean defining existence of d in d_unique
print(Iloc)
# [ True False False  True  True]

# indexes of d_unique that exists in d
print(idx)
# [4 4 3]

print(B[idx])
# [3 3 6]

print(A[Iloc])
# [3 3 6]

# These vectors will match
A[Iloc]==B[idx]

速度检查:

from ismember import ismember
from datetime import datetime

t1=[]
t2=[]
# Create some random vectors
ns = np.random.randint(10,10000,1000)

for n in ns:
    a_vec = np.random.randint(0,100,n)
    b_vec = np.random.randint(0,100,n)

    # Run stack version
    start = datetime.now()
    out1=ismember_stack(a_vec, b_vec)
    end = datetime.now()
    t1.append(end - start)

    # Run ismember
    start = datetime.now()
    out2=ismember(a_vec, b_vec)
    end = datetime.now()
    t2.append(end - start)


print(np.sum(t1))
# 0:00:07.778331

print(np.sum(t2))
# 0:00:04.609801

# %%
def ismember_stack(a, b):
    bind = {}
    for i, elt in enumerate(b):
        if elt not in bind:
            bind[elt] = i
    return [bind.get(itm, None) for itm in a]  # None can be replaced by any other "not in b" value

pypi的ismember函数快了将近2倍。

大向量,例如700000个元素:

from ismember import ismember
from datetime import datetime

A = np.random.randint(0,100,700000)
B = np.random.randint(0,100,700000)

# Lookup
start = datetime.now()
Iloc,idx = ismember(A, B)
end = datetime.now()

# Print time
print(end-start)
# 0:00:01.194801