python - 在 numpy 数组上并行化区域计算

Question

我尝试计算 numpy 数组上同一区域(相同值)的所有单元格的模式。我给你下面的代码示例。在这个例子中，顺序方法工作正常，但多进程方法什么都不做。我没有发现我的错误。

有人看到我的错误了吗？

我想并行化计算，因为我的实际数组是一个 10k * 10k 数组，有 100 万个区域。

import numpy as np
import scipy.stats as ss
import multiprocessing as mp

def zone_mode(i, a, b, output):
    to_extract = np.where(a == i)
    val = b[to_extract]
    output[to_extract] = ss.mode(val)[0][0]
    return output

def zone_mode0(i, a, b):
    to_extract = np.where(a == i)
    val = b[to_extract]
    output = ss.mode(val)[0][0]
    return output

np.random.seed(1)

zone = np.array([[1, 1, 1, 2, 3],
                 [1, 1, 2, 2, 3],
                 [4, 2, 2, 3, 3],
                 [4, 4, 5, 5, 3],
                 [4, 6, 6, 5, 5],
                 [6, 6, 6, 5, 5]])
values = np.random.randint(8, size=zone.shape)

output = np.zeros_like(zone).astype(np.float)

for i in np.unique(zone):
    output = zone_mode(i, zone, values, output)

# for multiprocessing    
zone0 = zone - 1

pool = mp.Pool(mp.cpu_count() - 1)
results = [pool.apply(zone_mode0, args=(u, zone0, values)) for u in np.unique(zone0)]
pool.close()
output = results[zone0]

Answer 1

对于数组中的正整数 - zone 和 values，我们可以使用 np.bincount。基本思想是我们将 zone 和 values 视为二维网格上的行和列。因此，可以将它们映射到它们的线性索引等效数。这些将用作与 np.bincount 合并求和的容器。它们的 argmax ID 将是模式编号。它们通过索引到 zone 映射回 zone-grid。

因此，解决方案是——

m = zone.max()+1
n = values.max()+1
ids = zone*n + values
c = np.bincount(ids.ravel(),minlength=m*n).reshape(-1,n).argmax(1)
out = c[zone]

---

对于稀疏数据(在输入数组中很好地分布整数)，我们可以查看稀疏矩阵以获得 argmax IDs c。因此，使用 SciPy 的稀疏矩阵 -

from scipy.sparse import coo_matrix

data = np.ones(zone.size,dtype=int)
r,c = zone.ravel(),values.ravel()
c = coo_matrix((data,(r,c))).argmax(1).A1

对于轻微的性能。提升，指定形状 -

c = coo_matrix((data,(r,c)),shape=(m,n)).argmax(1).A1

---

求解通用值

我们将使用 pandas.factorize，就像这样 -

import pandas as pd

ids,unq = pd.factorize(values.flat)
v = ids.reshape(values.shape)
# .. same steps as earlier with bincount, using v in place of values
out = unq[c[zone]]

请注意，对于并列情况，它会从 values 中挑选随机元素。如果您想选择第一个，请使用 pd.factorize(values.flat, sort=True)。