python:並列処理で得られる配列の扱いについて

Question

前提・実現したいこと

以下のコードを作成しました。

pDM = np.arrray([125., 256., 785., 34., 85., 73., 14., 3., 4.]
                [12., 26., 85., 3., 5., 713., 154., 32., 45.]
               [25., 26., 75., 134., 185., 273., 314., 33., 45.])
pGas = np.array([0.3, 4.5, 46.])
boxSize = 876.
def dist2(dx, dy, dz):
    return dx*dx + dy*dy + dz*dz

のような要素数の異なるpDMとpGasが存在します。

DM_adress = np.sqrt(dist2(pGas[1,0]-pDM[:,0], pGas[1,1]-pDM[:,1], pGas[1,2]-pDM[:,2], boxSize)) < (boxSize/512.0)
for k in range(pGas.shape[0])[1:]:
    DM_adress2 = np.sqrt(dist2(pGas[k,0]-pDM[:,0], pGas[k,1]-pDM[:,1], pGas[k,2]-pDM[:,2])) < (boxSize/512.0)
    DM_adress = np.vstack((DM_adress, DM_adress2))

(このとき，出来たDM_adress2を足し合わせる方法が分からなくて始めにk=0でDM_adressを作成して，それにfor文内でDM_adress2を足し合わせるという方法をとりました。)

並列計算の試行

実際はlen(pGas) > 10000, len(pDM) > 10^8 あり，この方法だと8時間以上回してもfor文の計算が終わりませんでした。
そこで並列処理を行おうとして以下のコードを作成しました。

from multiprocessing import Pool
def find_DMadress(k):
    found_DMad = np.sqrt(dist2(pGas[k,0]-pDM[:,0], pGas[k,1]-pDM[:,1], pGas[k,2]-pDM[:,2])) < (boxSize/512.0)
    return found_DMad

with Pool(processes=16) as p:
　　 DM_adress = p.map(find_DMadress, pGas)

最後のDM_adressの部分を書いていて手が止まってしまいました。
そこで質問は以下の通りです。
① 上の様に並列計算を行ったとき，found_DMadでは１次元配列が出来ますが、p.mapで出来るDM_adressには(len(pGas), len(pDM))の２次元配列ができるのでしょうか? 出来る為にはどのように変更すればよろしいでしょうか?

② 別の並列処理の方法を使用する等で、pGasとpDMの全ての行の組み合わせを行えるようにpGas, pDMの二つを引数にとって並列処理をすることは出来るのでしょうか?

①②それぞれ調べてもよく分からなかったので質問いたしました。
長文失礼いたしました。

よろしくお願いいたします。

補足情報（FW/ツールのバージョンなど）

python3

Comments

[hide5stm]

数秒で終る程度の小さいデータ量で計算するとどうでしょう

[jbpb0]

質問のコードを動かそうとしてもエラー出るし、質問を読んでもやりたい事がよく分からなかったのですが、下記で合ってますか？
・pDMとpGasには、それぞれたくさんの3次元座標 (x, y, z) が入ってる
・pDMとpGasの全ての組み合わせで距離を計算して、それが閾値よりも大きいか小さいかを判定した結果を計算したい (判定結果が格納された2次元配列を作りたい)

[Lily_1007]

はい、まさにそうです。閾値よるも大きい小さいでTrue, Falseで出力でも、
各pGasに対して閾値より小さい距離範囲にあるpDMの格納番号だけ取ってくるような出力のどちらかを得たいと思っております。

[jbpb0]

https://blog.shikoan.com/distance-without-for-loop/
を参考に、forを使わないコードを書いてみました
並列化とか考える前に、とりあえず、計算が合ってるか確認

import numpy as np

boxSize = 876.
th = boxSize / 512.0

pGas = np.array([
[0, 0, 0],
[0, 0, 1],
[0, 0, 2],
[0, 0, 3],
[0, 0, 4]
])

pDM = np.array([
[0, 0, 0],
[1, 1, 1],
[2, 2, 2]
])

all_distance = np.sqrt(np.sum((np.expand_dims(pGas, axis=1) - np.expand_dims(pDM, axis=0)) ** 2, axis=-1))
DM_adress = all_distance < th

print(all_distance)
print(DM_adress)

実行結果
>>> print(all_distance)
[[0. 1.73205081 3.46410162]
[1. 1.41421356 3. ]
[2. 1.73205081 2.82842712]
[3. 2.44948974 3. ]
[4. 3.31662479 3.46410162]]
>>> print(DM_adress)
[[ True False False]
[ True True False]
[False False False]
[False False False]
[False False False]]

欲しい結果の計算は、これで合ってますよね？

[Lily_1007]

はい、まさにその通りです。こちらを並列化したいと思っております。
(このときpGasだけを並列計算可能なのでしょうか?pDM. pGasを同時に非キス取っての並列計算も可能なのでしょうか.)

Answer 1

len(pGas) > 10000, len(pDM) > 10^8

len(pGas)：10000, len(pDM)：10^6 の場合の計算時間を下記コードで測ったら、当方のPCで150秒くらいでした

import numpy as np
import time

boxSize = 876.
th = boxSize / 512.0

pGas = np.random.rand(10000, 3).astype(np.float32)
pDM = np.random.rand(10**6, 3).astype(np.float32)

def DM_adress_calc(pGas, pDM):
  DM_adress_tmp = np.empty((len(pDM), len(pGas)), dtype=bool)
  for i in range(len(pDM)):
    DM_adress_tmp[i] = np.sqrt(np.sum((pGas - pDM[i])**2, axis=1)) < th
  return DM_adress_tmp.T

time1 = time.time()
DM_adress = DM_adress_calc(pGas, pDM)
time2 = time.time()
print(time2 - time1)

len(pGas)：10000, len(pDM)：10^8 の場合の計算時間を比例計算で見積もると、
150 * 10^8 / 10^6 / (60 * 60) ≒ 4時間
となりましたので、本当に4時間くらいで終わるか確認しようと、下記のコードを実行してみたら、
MemoryError: Unable to allocate 931. GiB for an array with shape (100000000, 10000) and data type bool
というエラーになってしまいました

import numpy as np
import time

boxSize = 876.
th = boxSize / 512.0

pGas = np.random.rand(10000, 3).astype(np.float32)
pDM = np.random.rand(10**8, 3).astype(np.float32)

def DM_adress_calc(pGas, pDM):
  DM_adress_tmp = np.empty((len(pDM), len(pGas)), dtype=bool)
  for i in range(len(pDM)):
    DM_adress_tmp[i] = np.sqrt(np.sum((pGas - pDM[i])**2, axis=1)) < th
  return DM_adress_tmp.T

time1 = time.time()
DM_adress = DM_adress_calc(pGas, pDM)
time2 = time.time()
print(time2 - time1)

bool型のサイズは1Byte単位なので、全ての組み合わせの結果を格納できる10000x10^8というサイズのbool型の配列のサイズは、
1000010^8/(10241024*1024) = 931GB
となってしまい、普通のPCではメモリーに収まりませんので、並列化で高速化するかどうかに関わらず、
・予めデータを分割しておいて、それぞれを分けて処理する
・ある程度処理したら、結果をファイルに書き出してメモリーから削除する
等の対策をして、計算した結果を格納する配列がメモリーに収まるサイズになるようにしないと、計算ができないと思います

たとえば、この回答の最初に挙げたコードは len(pDM)：10^8 を予め100等分しておいた場合に相当し、それならば計算できるので、それをデータを変えながら100回繰り返すようなコードを書けば、len(pGas)：10000, len(pDM)：10^8 の全ての組み合わせを計算することができます

そのような方法で計算することが前提になりますが、numba を用いて高速化して、len(pGas)：10000, len(pDM)：10^6 の場合の計算時間を下記コードで測ったら、当方のPCで25秒くらいでした

import numpy as np
from numba import jit, prange
from numba.types import bool_, float32
import time

boxSize = 876.
th = boxSize / 512.0

pGas = np.random.rand(10000, 3).astype(np.float32)
pDM = np.random.rand(10**6, 3).astype(np.float32)

@jit(bool_[:, :](float32[:, :], float32[:, :]), nopython=True, fastmath=True, nogil=True, parallel=True)
def DM_adress_calcj(pGas, pDM):
  DM_adress_tmp = np.empty((len(pDM), len(pGas)), dtype=bool_)
  for i in prange(len(pDM)):
    DM_adress_tmp[i] = np.sqrt(np.sum((pGas - pDM[i])**2, axis=1)) < th
  return DM_adress_tmp.T

time1 = time.time()
DM_adress = DM_adress_calcj(pGas, pDM)
time2 = time.time()
print(time2 - time1)

len(pGas)：10000, len(pDM)：10^8 の場合の計算時間を比例計算で見積もると、
25 * 10^8 / 10^6 / 60 ≒ 40分
となり、numba を用いない場合の予想約4時間の1/6で済む予想です

numba を使える環境をお持ちでしたら、簡単に高速化できるので、お勧めです
Anaconda版Pythonなら、普通にインストールするだけで numba が使える状態になっていたと思います