使用matplotlib plot_surface遇矩阵求逆报错的解决咨询

问题描述 #

使用matplotlib的plot_surface绘制3D曲面时，自定义的f(x,y)函数因矩阵求逆操作触发错误：Last 2 dimensions of the array must be square，但用plot绘制线框曲面时无异常，需解决该矩阵求逆问题以正常绘制曲面。

原代码 #

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
from sympy import *

fig = plt.figure()
ax = Axes3D(fig)
x = np.arange(52, 100, 4)
y = np.arange(1.5, 2, 0.05)
X,Y = np.meshgrid(x,y)

def f(x,y):
    L = 0.53*y
    l_2 = 0.2*y
    x_s = 0.004 
    m_1 = 0.19*x
    m_2 = 0.62*x
    l_1 = 0.3167*y
    
    R = 0.45
    r = 0.41
    h = 0.38
    b = 0.2
    d = 0.12
    M_0 = 2.2

    k = 5
    g = 9.8

    Q = x*g*x_s
    I_1 = M_0*d**2
    I_2 = m_1*L**2/3
    I_3 = m_2*(b**2+l_2**2)/3
    A_0 = M_0*(R-r)**2+2*m_1*((R-h)**2+b**2)+m_2*(R-h+l_2)**2+I_1+I_3
    B_0 = 2*m_1*l_1*(R-h)+m_2*L*(R-h+l_2)
    C_0 = 2*m_1*l_1**2+m_2*L**2+2*I_2
    w = g*(M_0*r+2*m_1*h-m_2*(l_2-h))
    v = -g*(2*m_1*l_1+m_2*L)
    
    
    M = np.array([[A_0, B_0], [B_0, C_0]])
    N = np.array([[k, -k], [-k, k]])
    K = np.array([[w, 0], [0,v]])
    D = np.array([-1,1])

    M_1 = np.linalg.inv(M) #error

    N_0 = np.dot(M_1, N)
    K_0 = np.dot(M_1, K)
    D_0 = np.dot(M_1, D)
    D_0 = D_0.reshape(-1,1)

    F_0 = np.concatenate ([eye(2), zeros(2)], axis = 1)
    F = np.concatenate([-N_0, -K_0], axis = 1)
    F = np.concatenate([F, F_0])
    F = np.array(F, dtype = float)
    zer = np.array([[0],[0]])
    H = np.concatenate([D_0, zer])
    t,T = np.linalg.eig(F)

    for i in range(4):
        if re(t[i]) > 0:
            j = i
    
    T_0 = np.linalg.inv(T)
    P = np.dot(T_0,H)
    p = P[j,0]
    s = t[j]
    s = float(s)
    return Q*abs(p)*180/(pi*s*abs(T_0[j,2]))

ax.plot_surface(X,Y,f(X,Y))
plt.show()

错误原因 #

调用f(X,Y)时，X和Y是二维网格数组（示例中shape为(10,12)），函数内的A_0、B_0、C_0等变量都会变成同shape的二维数组。此时构造的M = np.array([[A_0, B_0], [B_0, C_0]])是一个4维数组，shape为(2,2,10,12)，而np.linalg.inv要求最后两维是方阵，但当前维度顺序错误，导致无法直接求逆。而plot线框绘制时可能是逐点处理标量输入，因此不会触发该错误。

解决方案 #

方案1：向量化批量处理矩阵操作（高效） #

调整矩阵构造方式，让M的shape为(n,m,2,2)，使np.linalg.inv能批量处理每个网格点的2x2矩阵，同时用np.matmul或@完成批量矩阵乘法：

修改后的完整代码：

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = Axes3D(fig)
x = np.arange(52, 100, 4)
y = np.arange(1.5, 2, 0.05)
X,Y = np.meshgrid(x,y)

def f(X,Y):
    L = 0.53*Y
    l_2 = 0.2*Y
    x_s = 0.004 
    m_1 = 0.19*X
    m_2 = 0.62*X
    l_1 = 0.3167*Y
    
    R = 0.45
    r = 0.41
    h = 0.38
    b = 0.2
    d = 0.12
    M_0 = 2.2

    k = 5
    g = 9.8

    Q = X*g*x_s
    I_1 = M_0*d**2
    I_2 = m_1*L**2/3
    I_3 = m_2*(b**2+l_2**2)/3
    A_0 = M_0*(R-r)**2+2*m_1*((R-h)**2+b**2)+m_2*(R-h+l_2)**2+I_1+I_3
    B_0 = 2*m_1*l_1*(R-h)+m_2*L*(R-h+l_2)
    C_0 = 2*m_1*l_1**2+m_2*L**2+2*I_2
    w = g*(M_0*r+2*m_1*h-m_2*(l_2-h))
    v = -g*(2*m_1*l_1+m_2*L)
    
    # 构造shape为(n,m,2,2)的批量矩阵
    M = np.stack([np.stack([A_0, B_0], axis=-1), 
                  np.stack([B_0, C_0], axis=-1)], axis=-2)
    # 批量求逆
    M_1 = np.linalg.inv(M)
    
    # 扩展N的维度以支持批量乘法
    N = np.array([[k, -k], [-k, k]])[np.newaxis, np.newaxis, :, :]
    N_0 = np.matmul(M_1, N)
    
    # 构造shape为(n,m,2,2)的K矩阵
    K = np.stack([np.stack([w, np.zeros_like(w)], axis=-1), 
                  np.stack([np.zeros_like(v), v], axis=-1)], axis=-2)
    K_0 = np.matmul(M_1, K)
    
    # 扩展D的维度以支持批量乘法
    D = np.array([-1,1])[np.newaxis, np.newaxis, :, np.newaxis]
    D_0 = np.matmul(M_1, D).squeeze(-1)
    
    # 构造批量F矩阵（shape为(n,m,4,4)）
    eye_2 = np.eye(2)[np.newaxis, np.newaxis, :, :]
    zeros_2x2 = np.zeros((2,2))[np.newaxis, np.newaxis, :, :]
    F_0 = np.concatenate([eye_2, zeros_2x2], axis=-1)
    F_top = np.concatenate([-N_0, -K_0], axis=-1)
    F = np.concatenate([F_top, F_0], axis=-2)
    
    # 批量计算特征值和特征向量
    t, T = np.linalg.eig(F)
    
    # 找到每个网格点实部大于0的特征值索引
    j = np.argmax(np.real(t) > 0, axis=-1)
    # 调整索引形状以适配后续索引
    j = j[..., np.newaxis, np.newaxis]
    
    # 批量求逆T矩阵
    T_0 = np.linalg.inv(T)
    
    # 构造H矩阵（shape为(n,m,4,1)）
    zer = np.zeros_like(D_0)[..., np.newaxis]
    H = np.concatenate([D_0[..., np.newaxis], zer], axis=-2)
    
    # 批量计算P
    P = np.matmul(T_0, H)
    
    # 提取对应索引的p和s
    p = np.take_along_axis(P, j, axis=-2).squeeze((-2,-1))
    s = np.take_along_axis(t, j[...,0], axis=-1).squeeze(-1)
    s = s.astype(float)
    
    # 提取T_0[j,2]的值
    T0_j2 = np.take_along_axis(T_0[...,2], j[...,0], axis=-1).squeeze(-1)
    
    return Q*np.abs(p)*180/(np.pi*s*np.abs(T0_j2))

Z = f(X,Y)
ax.plot_surface(X,Y,Z)
plt.show()

方案2：用np.vectorize包装标量函数（简单但低效） #

将原函数修改为处理单个标量x和y，再用np.vectorize包装，使其支持数组输入：

修改后的完整代码：

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure()
ax = Axes3D(fig)
x = np.arange(52, 100, 4)
y = np.arange(1.5, 2, 0.05)
X,Y = np.meshgrid(x,y)

def f_scalar(x,y):
    L = 0.53*y
    l_2 = 0.2*y
    x_s = 0.004 
    m_1 = 0.19*x
    m_2 = 0.62*x
    l_1 = 0.3167*y
    
    R = 0.45
    r = 0.41
    h = 0.38
    b = 0.2
    d = 0.12
    M_0 = 2.2

    k = 5
    g = 9.8

    Q = x*g*x_s
    I_1 = M_0*d**2
    I_2 = m_1*L**2/3
    I_3 = m_2*(b**2+l_2**2)/3
    A_0 = M_0*(R-r)**2+2*m_1*((R-h)**2+b**2)+m_2*(R-h+l_2)**2+I_1+I_3
    B_0 = 2*m_1*l_1*(R-h)+m_2*L*(R-h+l_2)
    C_0 = 2*m_1*l_1**2+m_2*L**2+2*I_2
    w = g*(M_0*r+2*m_1*h-m_2*(l_2-h))
    v = -g*(2*m_1*l_1+m_2*L)
    
    M = np.array([[A_0, B_0], [B_0, C_0]])
    N = np.array([[k, -k], [-k, k]])
    K = np.array([[w, 0], [0,v]])
    D = np.array([-1,1])

    M_1 = np.linalg.inv(M)

    N_0 = np.dot(M_1, N)
    K_0 = np.dot(M_1, K)
    D_0 = np.dot(M_1, D)
    D_0 = D_0.reshape(-1,1)

    F_0 = np.concatenate ([np.eye(2), np.zeros((2,2))], axis = 1)
    F = np.concatenate([-N_0, -K_0], axis = 1)
    F = np.concatenate([F, F_0])
    F = np.array(F, dtype = float)
    zer = np.array([[0],[0]])
    H = np.concatenate([D_0, zer])
    t,T = np.linalg.eig(F)

    # 找到第一个实部大于0的特征值索引
    j = np.where(np.real(t) > 0)[0][0]
    
    T_0 = np.linalg.inv(T)
    P = np.dot(T_0,H)
    p = P[j,0]
    s = t[j]
    s = float(s)
    return Q*abs(p)*180/(np.pi*s*abs(T_0[j,2]))

# 向量化函数，使其支持数组输入
f = np.vectorize(f_scalar)

ax.plot_surface(X,Y,f(X,Y))
plt.show()

额外注意事项 #

• 移除了不必要的sympy导入，避免和numpy的eye、zeros函数混淆；
• 方案1适合大数据量场景，效率更高；方案2实现简单，但本质是循环处理每个点，数据量大时速度较慢。

内容的提问来源于stack exchange，提问作者Luna Lovegood