関数の引数の個数と返り値の個数をif文を使わずに変更したい。

Question

前提・実現したいこと

リストgridの要素としてnumpyのnumpy.meshgridの返り値を代入しているのですが、dimensionの値によって、引数の数と返り値の数を変更するようにしたいです。
現状ではif文で分岐させていますが、見づらいのとdimension == 100等の場合、対応が難しいです。

該当のソースコード

axis_grid = []  
for i in range(5):
    axis_grid.append(np.linspace(-15, 15, 10))

dimension = len(axis_grid)

grid = []
for i in range(dimension):
    grid.append([])
if dimension == 2:
    grid[0], grid[1] = np.meshgrid(axis_grid[0], axis_grid[1])
elif dimension == 3:
    grid[0], grid[1], grid[2] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2])
elif dimension == 4:
    grid[0], grid[1], grid[2], grid[3] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3])
elif dimension == 5:
    grid[0], grid[1], grid[2], grid[3], grid[4] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3], axis_grid[4])
elif dimension == 6:
    grid[0], grid[1], grid[2], grid[3], grid[4], grid[5] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3], axis_grid[4], axis_grid[5])
elif dimension == 7:
    grid[0], grid[1], grid[2], grid[3], grid[4], grid[5], grid[6] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3], axis_grid[4], axis_grid[5], axis_grid[6])
elif dimension == 8:
    grid[0], grid[1], grid[2], grid[3], grid[4], grid[5], grid[6], grid[7] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3], axis_grid[4], axis_grid[5], axis_grid[6], axis_grid[7])
elif dimension == 9:
    grid[0], grid[1], grid[2], grid[3], grid[4], grid[5], grid[6], grid[7], grid[8] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3], axis_grid[4], axis_grid[5], axis_grid[6], axis_grid[7], axis_grid[8])
elif dimension == 10:
    grid[0], grid[1], grid[2], grid[3], grid[4], grid[5], grid[6], grid[7], grid[8], grid[9] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3], axis_grid[4], axis_grid[5], axis_grid[6], axis_grid[7], axis_grid[8], axis_grid[9])
elif dimension == 11:
    grid[0], grid[1], grid[2], grid[3], grid[4], grid[5], grid[6], grid[7], grid[8], grid[9], grid[10] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3], axis_grid[4], axis_grid[5], axis_grid[6], axis_grid[7], axis_grid[8], axis_grid[9], axis_grid[10])
elif dimension == 12:
    grid[0], grid[1], grid[2], grid[3], grid[4], grid[5], grid[6], grid[7], grid[8], grid[9], grid[10], grid[11] = np.meshgrid(axis_grid[0], axis_grid[1], axis_grid[2], axis_grid[3], axis_grid[4], axis_grid[5], axis_grid[6], axis_grid[7], axis_grid[8], axis_grid[9], axis_grid[10], axis_grid[11])

Comments

[tetsunosuke]

回答というわけではないのですがぱっと検索してみたら
https://www.it-swarm-ja.tech/ja/python/3d%e3%81%aenumpy-meshgrid/968575565/

こちらのmeshgrid2を使うことで可変長引数にして、meshgrid2(axis_grid) と呼び出すことができるように見えました。どうかな...

Answer 1

リストを分解して、引数に渡したい場合はアンパックを使います。

4. その他の制御フローツール — Python 3.8.4 ドキュメント

def f(a, b, c):
  print(a, b, c)

vals = [1, 2, 3]
f(*vals)  # 1 2 3

これを使って上記コードを簡略化すると、以下のようになります。

import numpy as np

dims = 3  # 次元数
vals = np.linspace(-15, 15, 5)  # 範囲

# [vals] * dims は [vals, vals, ..., vals] という配列を作るのに利用
args = [vals] * dims
grid = np.meshgrid(*args)
print(grid)

ちなみに次元を増やすと必要なメモリ量は指数関数的に増えますので、meshgrid() は普通 dims=2, 3 あたりで使うものであり、高次元の配列を作るということは基本的にないと思います。

len(vals) = 5 の場合

	dim	number of elements	bytes (float64)
0	1	10	80.0B
1	2	200	1.6KiB
2	3	3000	23.4KiB
3	4	40000	312.5KiB
4	5	500000	3.8MiB
5	6	6000000	45.8MiB
6	7	70000000	534.1MiB
7	8	800000000	6.0GiB
8	9	9000000000	67.1GiB
9	10	100000000000	745.1GiB