numpy.matmul对dtype=object运算规则及ABCD光线传输计算疑问

问题背景 #

我在使用numpy.matmul处理dtype=object类型的数组时存在困惑，希望能了解下述代码的运行逻辑，尤其是最后一步得到的total变量的具体内容。我尝试用该方法替代循环迭代实现矩阵乘法，不确定该实现是否正确。
我正在进行同时可变距离z、输入光束角g_in的ABCD光线传输模拟计算，计算结果将用于生成两张热力图：纵轴为z，横轴为g_in，两张图的色条值分别对应total返回结果的第一个、第二个元素。

代码实现 #

import numpy as np
import matplotlib.pyplot as plt

z = np.linspace(0, 10, 11) # 传输距离
g_in = np.linspace(-5, 5, 11) # 输入光束角度
f_0, f_1 = 1, 1 # 光学系统常数

A = np.array([[1, z], [0, 1]], dtype = object)
B = np.array([[1, 0], [-1/(f_0), 1]], dtype = object)
C = np.array([[1, 2*f_0], [0, 1]], dtype = object)
D = np.array([[1, 0], [-1/(f_0), 1]], dtype = object)
E = np.array([[1, z], [0, 1]], dtype = object)
F = np.array([[1, 0], [-1/(f_1), 1]], dtype = object)
G = np.array([[1, f_1], [0, 1]], dtype = object)

H = np.matmul(G,F)
I = np.matmul(H,E)
J = np.matmul(I,D)
K = np.matmul(J,C)
L = np.matmul(K,B)

M = np.matmul(L, A)

print('Matrix M=',M)

col1 = np.empty((2, 1),dtype=object)
col1[:, 0] = [0, g_in]

print('Matrix col1[:,0]=',col1[:,0])

total = np.matmul(M, col1[:,0])

print('Matrix total=',total)

y_out = np.transpose(total[0].tolist())
g_out = np.transpose(total[1].tolist())

y_out_ = np.expand_dims(y_out, axis=0)
g_out_ = np.expand_dims(g_out, axis=0)

fig, ax1 = plt.subplots(nrows=1,
ncols=1,sharex=True,sharey=True, figsize=(8, 6))
f1=ax1.imshow(y_out_, extent=  
[g_in.min(),g_in.max(),z.min(),z.max()],    
vmin=y_out_.min(),vmax=y_out_.max(), aspect="auto",
cmap='YlGnBu')
cb1=fig.colorbar(f1,orientation='vertical')
cb1.set_label(r'$y_{out}$',size=15)
ax1.set_ylabel(r'z', fontsize=20)
ax1.tick_params(axis='both', which='major', labelsize=20)
ax1.tick_params(axis='both', which='minor', labelsize=20)
ax1.autoscale(tight=True)

fig, ax2 = plt.subplots(nrows=1, ncols=1,sharex=True,figsize=(8, 6))
f2=ax2.imshow(g_out_, extent= 
[g_in.min(),g_in.max(),z.min(),z.max()],   
vmin=g_out_.min(),vmax=g_out_.max(), aspect="auto",
cmap='YlGnBu')
cb2=fig.colorbar(f2,orientation='vertical')
cb2.set_label(r'$g_{out}$',size=15)
ax2.set_xlabel(r' Angle, $θ_{in}$', fontsize=20)
ax2.set_ylabel(r'z', fontsize=20)
ax2.tick_params(axis='both', which='major', labelsize=20)
ax2.tick_params(axis='both', which='minor', labelsize=20)
ax2.autoscale(tight=True)
plt.show()

运行结果 #

1. object类型数组的运算逻辑
   你用到的矩阵A、E中包含长度为11的z数组，因此使用dtype=object可以直接将数组作为矩阵的单个元素存储。numpy.matmul处理object类型矩阵时，会自动对元素执行对应算术运算，等效于隐式遍历所有z取值完成矩阵乘法，不需要手动写循环。

2. total变量的具体内容

• 经过级联乘法得到的M是2×2的object数组，每个元素都是长度为11的数组，对应11个z取值下的ABCD矩阵元素值
• 输入向量col1[:,0]是长度为2的object数组，第一个元素是标量0（输入光线初始位置），第二个是长度为11的g_in数组（输入光线角度）
• total = np.matmul(M, col1[:,0])是长度为2的object数组：
  • total[0]：11×11的二维数组，行对应z的取值，列对应g_in的取值，每个位置的数值为对应参数下的输出光线位置y_out
  • total[1]：11×11的二维数组，每个位置的数值为对应参数下的输出光线角度g_out

3. 实现正确性判断
   该实现逻辑完全正确，本质是向量化的ABCD矩阵运算，和嵌套循环遍历z、g_in计算的结果完全一致，且运算效率远高于手动实现的循环。

4. 代码小修正
   原代码绘图部分使用了未定义的theta_in、z_f变量，替换为前面定义的g_in、z即可正常运行。

内容的提问来源于stack exchange，提问作者Christofer Iacovou