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

Question

我尝试在 numpy 数组上的同一区域（相同值）的所有单元格上计算模式。我在下面给你一个代码示例。在此示例中，顺序方法工作正常，但多处理方法无济于事。我没有发现我的错误。

有人看到我的错误吗？

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

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 进行合并的bin。他们的argmax ID 将是模式编号。它们被映射回 zone-grid，索引为 zone。

因此，解决方案是 -

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]

---

对于稀疏数据（输入数组中分布良好的整数），我们可以查看 sparse-matrix 以获取 argmax ID 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

---

求解泛型values

我们将使用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]]

请注意，对于 tie-cases，它会从 values 中选择随机元素。如果要选择第一个，请使用pd.factorize(values.flat, sort=True)。