Numpy:填充数组以矢量化嵌套求和

Question

我想在Python中对嵌套和进行矢量化，以加快进程。目前我已经嵌套了for循环。你知道吗

    for ja in np.arange(0,Na):
        for jb in np.arange(0,Nb):
            for ma in np.arange(-ja,ja+1):
            ...

最终结果是2x2矩阵的和，每个矩阵的条目取决于ja、jb、ma、mb的值。 矩阵如下所示：

[[f11(ja,jb,ma,mb),f12(ja,jb,ma,mb)],
 [f21(ja,jb,ma,mb),f22(ja,jb,ma,mb)]]

其中fij是函数。这些函数可以应用于数组，因为它们将逐个元素（指数、平方根、三角函数等）工作。我可以创建如下数组：

ja=[0,0,0,1,1,1,2,2,2,3,3,3]
jb=[0,1,2,3,0,1,2,3,0,1,2,3]

通过使用

range_a = np.arange(0,Na/2+1)
range_b = np.arange(0,Nb/2+1)
ja = np.tile(a_range,Nb/2+1)
jb = np.repeat(b_range,Na/2+1)

但我的问题是创建数组，这样对于上面j中的每个值，我们都有m结构（从-j到j）：

ma=[0,0,0,-1,-1,-1,0,0,0,1,1,1,-2,-2,-2,...]
mb=[0,-1,0,1,-2,-1,0,1,2,...]

我有困难，使这些m阵列虽然！每次a-j，…，j结构重复时，它都有不同的长度，所以我不能使用像numpy.tile和numpy.repeat这样的函数。那么，有什么办法吗？你知道吗

我的进一步意图可能与此相关：我希望能够用零填充它们，并构造2*长度的矩阵，以便只填充一个条目（每个向量需要4个矩阵，总共4个向量）。然后我可以将函数应用到并添加这些upp。为了计算2*2矩阵的和，我将用一个长度为*4的矩阵做点积。结果是2*2。也许有更好的策略存在？我认为这种情况以前可能发生过，因为它在物理学中有一个常见的应用（密度算符上的迹），但我没有找到它。你知道吗