Numpy 矩阵操作

Question

我想计算所有的 i 和 j 的以下值：

M_ki = Sum[A_ij - A_ik - A_kj + A_kk, 1 <= j <= n]

我该如何使用 Numpy（Python）来做到这一点，而不使用显式的循环呢？

谢谢！

Answer 1

这里有一个解决这类问题的一般策略。

首先，写一个小脚本，里面有两个不同的函数，循环部分要写得很清楚，最后要有一个测试，确保这两个函数的结果完全一样：

import numpy as np
from numpy import newaxis

def explicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    for k in range(n):
        for i in range(n):
            for j in range(n):
                m[k,i] += a[i,j] - a[i,k] - a[k,j] + a[k,k]
    return m

def implicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    for k in range(n):
        for i in range(n):
            for j in range(n):
                m[k,i] += a[i,j] - a[i,k] - a[k,j] + a[k,k]
    return m

a = np.random.randn(10,10)
assert np.allclose(explicit(a), implicit(a), atol=1e-10, rtol=0.)

接下来，逐步将这个函数进行向量化，修改 implicit，每一步都运行脚本，确保它们的结果依然相同：

第一步

def implicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    for k in range(n):
        for i in range(n):
            m[k,i] = (a[i,:] - a[k,:]).sum() - n*a[i,k] + n*a[k,k]
    return m

第二步

def implicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    m = - n*a.T + n*np.diag(a)[:,newaxis]
    for k in range(n):
        for i in range(n):
            m[k,i] += (a[i,:] - a[k,:]).sum()
    return m

第三步

def implicit(a):
    n = a.shape[0]
    m = np.zeros_like(a)
    m = - n*a.T + n*np.diag(a)[:,newaxis]
    m += (a.T[newaxis,...] - a[...,newaxis]).sum(1)
    return m

好了！最后一个版本里没有循环。要将这类方程向量化，使用广播是个不错的选择！

注意：确保 explicit 是你想要向量化的方程。我不确定那些不依赖于 j 的项是否也应该被求和。

Numpy 矩阵操作

1 个回答

撰写回答