numpy でこの複雑な行列をどうやって効率的に作れますか

Question

前提・実現したいこと

画像のような構造の行列をPythonで効率的（高速で）作りたいです（h(i,j)等はi,jの関数）。for文を使いまくれば作れるとは思うのですが、Pythonらしく、numpy等を使ってforをなるべく使わずに作りたいです。numpyだと和の範囲が行列要素に依存しているような行列はうまくかけないのでしょうか、、、

（補足）
現在の僕のプロジェクトの計算の律速がこの行列作成となっているみたいなので、なるべく高速にこの行列を作る方法を御提示いただけると幸いです。具体的にはこのf(i,j,k)の関数系を少しづつ変えた行列を大量に作らなくては行けません。和の範囲の関数は変わりません。
![]

上の行列の具体例：

発生している問題・エラーメッセージ

少し具体化して例えばg(i,j) = j, h(i,j) = 7-i, f(i,j,k) = sin(i+j+k)として、添字は -2 <= i,j <= 2　という条件でこの行列を作るために、以下のようなコードを書くと、次のようなエラーが表示されます。ソースコードのnp.arange(i,7-j)のiとjには関数が呼び出された時にnumpy arrayが代入されているからだとはわかっているのですが、これをどうすればいいのかわかりません。ここが、和の範囲が行列要素に依存している部分です。

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

該当のソースコード

def test_func(i,j):
    ini_array = np.arange(j,7-i)
    return np.sin(i+j+ini_array).sum()

ii, jj = np.mgrid[-2:3,-2:3]
M_matrix = test_func(ii,jj)

試したこと

np.arrange(i,7-j)の部分をfor文で作っても結局は同じ問題に、、、（forの範囲があとで代入するnumpyに依存している）

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

python 3.7
import numpy as np

Answer 1

「numpy等を使ってforをなるべく使わずに作りたい」という部分が書き方や読みやすさに関する期待だけであって、numpy を使うことによる高速化を期待しないのであれば、下記のように numpy.vectorize 関数を使うことで、実現可能です。

def test_func(i,j):
    ini_array = np.arange(j,7-i)
    return np.sin(i+j+ini_array).sum()

test_func_vectorized=np.vectorize(test_func)

ii, jj = np.mgrid[-2:3,-2:3]
M_matrix = test_func_vectorized(ii,jj)

Comments

[Saikohage]

行列が大きい、かつなんども作らないと行けないため、numpyを使い高速化したかったのですが、vectorizeでひとまずはすっきりしたコードになることがわかり大変参考になりました。
ご回答ありがとうございました。

Answer 2

興味を持ったのでちょっと試してみました。

まずは、基本的なところで動作確認

import math as m

def g(i,j):
    return j

def h(i,j):
    return 7-i

def f(i,j,k):
    return m.sin(i+j+k)

def main1():
    ii = np.arange(-2,3)
    jj = np.arange(-2,3)

    frm = g(ii,jj)
    print("frm: ", frm)
    to = h(ii,jj)
    print("to: ", to)

    for j, t in enumerate(to):
        for i, fr in enumerate(frm):
            func = lambda k: f(i,j,k)
            # ar = np.arange(fr, t+1)
            ar = sum(map(func, range(fr, t+1)))
            print(f"{i}:{j} {ar}")

main1()

動作することはできました。これを効率的にできるかですが、計算時に
i,jが必要となるのが厄介です。

def main2():

    ii, jj = np.mgrid[-2:3,-2:3]
    frm = g(ii,jj)
    to  = h(ii,jj)

    def fn(x):
        return x
    it = np.nditer(frm, flags=["multi_index"])
    with it:
        while not it.finished:
            mi = it.multi_index
            ff = frm[mi]
            tt = to[mi]
            func = lambda k: f(mi[0],mi[1],k)
            ar = sum(map(func, range(ff, tt+1)))
            print(mi[0], mi[1], ar)
            it.iternext()

こんな感じでできますが、もう少しなんとかしたいところ。
計算時のi,jをあらかじめ用意するようにします。

def fromto(ii,jj):
    return [[(g(vi[j],vj), h(vi[j],vj), i, j) for j,vj in enumerate(jj[i])] for i,vi in enumerate(ii)]

def main3():
    ii, jj = np.mgrid[-2:3,-2:3]
    ft = fromto(ii,jj)

    def func(v):
        ff = lambda k: f(v[2],v[3],k)
        ar = sum(map(ff, range(v[0], v[1]+1)))
        return ar

    l = list(map(lambda x: list(map(func,x)), ft))
    print(l)

fromto関数で、計算に必要な情報をあらかじめ準備して、listに入れてしまいます。
結局mapで再度繰り返していますが、スッキリした感じがしますがいかがでしょうか。

Comments

[Saikohage]

非常に示唆に富んだご回答ありがとうございました。補足でもあるように、行列作成の速度をあげることに苦心しています。ご提示いただいた方法を試してみたいと思います。

Answer 3

解決しました。いろいろなアイディアをご教示いただいた皆さま、どうもありがとうございました。

F_{i,j,k} = f(i,j,k) となる３次元の配列Fと、(Mask)_{i,j,k} = True if k > g(i,j) and k < h(i,j)となるようなマスク配列Maskを用意し、F*Mask してから、kについて和をとることで実現しました。

例えばg(i,j) = j, h(i,j) = 7-i, f(i,j,k) = sin(i+j+k)として、添字は -2 <= i,j <= 2　という条件では以下のようなコードとなります。

#kkの範囲はg(i,j)の最小値以上、h(i,j)の最大値以下
ii, jj, kk = np.mgrid[-2:3, -2:3, -2:8]

# k > g(i,j)
frm = kk > jj

# k < h(i,j)
to = kk < 7 - ii 

#マスク配列
k_mask = np.logical_and(frm,to)

def f(i,j,k):
    return np.sin(i+j+k)

M_before_mask = f(ii,jj,kk)

#完成した行列
M_matrix = (M_before_mask*k_mask).sum(axis=2)
コード