Python:返回相同形状数组的数组函数

时间:2016-07-19 21:25:18

标签: python arrays function numpy

注意:我正在使用numpy

import numpy as np

给定4个相同(但任意)形状的数组,我试图编写一个函数,从数组的每个对应元素形成2x2矩阵,找到特征值,并返回两个与原始四个相同形状的数组,其元素为特征值(即生成的数组与输入的形状相同,array1保存所有第一个特征值,array2保存所有第二个特征值)。

我尝试了以下操作,但不出所料,它给了我一个错误,表明数组不是正方形。

temp = np.linalg.eig([[m1, m2],[m3, m4]])[0]

我想我可以制作一个相同形状的空temp变量,

temp = np.zeros_like(m1)

并遍历原始数组的每个元素并重复该过程。我的问题是我想要对任意形状的数组进行推广(不必是一维)。我猜想找到数组的形状并设计循环来遍历每个元素并不是一个非常好的方法。我该如何有效地做到这一点?

1 个答案:

答案 0 :(得分:1)

构造一个2x2x ...数组:

temp = np.array([[m1, m2], [m3, m4]])

将前两个维度移动到... x2x2数组的末尾:

for _ in range(2):
    temp = np.rollaxis(temp, 0, temp.ndim)

调用np.linalg.eigvals(广播)获取... x2特征值数组:

eigvals = np.linalg.eigvals(temp)

将其拆分为第一个特征值数组和第二个特征值数组:

eigvals1, eigvals2 = eigvals[..., 0], eigvals[..., 1]