[Q&A] pythonのforループの高速化

pythonのforループの高速化

解決したいこと

私が以前書いたQiitaの記事の中にあるpythonスクリプトについてです．この記事内のスクリプトの中に

~(中略)~

for ix in range(cell_dims[0]+1):
    for iy in range(cell_dims[1]+1):
        for iz in range(cell_dims[2]+1):
            x = ix * mesh_pitch[0] + mins[0]
            y = iy * mesh_pitch[1] + mins[1]
            z = iz * mesh_pitch[2] + mins[2]
            points.InsertNextPoint(x, y, z)
structured_base_mesh = vtk.vtkStructuredGrid()

~(中略)~

という記述があるのですが，pythonの３重ループということもあり，問題によっては非常に時間がかかってしまうという問題があります．
できるだけここを高速にしたい場合，どのような手段がありますでしょうか？

pybindなどは使わず，できるだけpythonで完結させたいです．

ご教授の程宜しくおねがいします．

2Answer

@shiracamus posted at 2021-11-07

itertoolsを使う
事前に変数代入しておく
変数代入をなくす

といったところでしょうか。

from itertools import product
nx, ny, nz = cell_dims
px, py, pz = mesh_pitch
mx, my, mz = mins
for ix, iy, iz in product(range(nx + 1), range(ny + 1), range(nz + 1)):
    points.InsertNextPoint(ix * px + mx,
                           iy * py + my,
                           iz * pz + mz)

関数にして変数辞書を小さくすれば、多少変わるかもしれません。

from itertools import product

def make_points(nx, ny, nz, px, py, pz, mx, my, mz):
    points = vtk.vtkPoints()
    for ix, iy, iz in product(range(nx + 1), range(ny + 1), range(nz + 1)):
        points.InsertNextPoint(ix * mx + mx, iy * my + my, iz * mz + mz)
    return points

points = make_points(*cell_dims, *mesh_pitch, *mins)

Comments

@dwatanabee
Questioner
採用させていただきました．
ありがとうございます．
@shiracamus
vtk を知らなくて恐縮ですが、numpy を使えるなら、ループで InsertNextPoint するのではなく SetData で一括登録できたりしないでしょうか？
検索してみたところ、points.SetData(numpy_support.numpy_to_vtk(vertices)) のような書き方が見つかります。
https://stackoverflow.com/questions/66341661/how-to-fasten-vtk-addpoint-loop-python
@dwatanabee
Questioner
setDatの存在は知りませんでした．．．ありがとうございます．
@_y_sさんのコメントと組み合わせたら実装することが出来ました．
概して，元のコードに対して3/5程度の計算時間になったみたいです．
下の記事は後で更新します．
https://qiita.com/dwatanabee/items/b938b46ef911a43b92ba

@_y_s posted at 2021-11-07

numpyを使えるならnumpyで先に座標を作ってしまうのもありかと思います。

import numpy as np
px, py, pz = mesh_pitch
mx, my, mz = cell_dims * mesh_pitch #max
coords = np.stack(np.mgrid[:mx:px, :my:py, :mz:pz], -1).reshape(-1, 3) + mins
for c in coords:
    points.InsertNextPoint(*c)

Comments

@dwatanabee
Questioner
ありがとうございます．
numpyのほうが早そうですが，そのまま実装するとエラーになるみたいなので，もう少し試行錯誤してみたいと思います．
@_y_s
説明不足ですみません。mesh_pitchやcell_dimsをndarrayにしてもらえれば大丈夫だと思います。

例：`cell_dims = np.array([mesh_size, mesh_size, mesh_size])`
@dwatanabee
Questioner
ありがとうございます。
書き直したらうまくいきました。
速さはitertoolsとトントンくらいでした。

pythonのforループの高速化

解決したいこと

2Answer

Comments

Comments

Your answer might help someone💌

Popular Questions