基于numpy构造满足B@X.ravel()=(A.T@X+X@A).ravel()的矩阵B

我对numpy并不陌生，但这个问题的解法我摸索了很久仍未找到，因此在此提问。
给定等式 Y = A.T @ X + X @ A，其中A和X均为同阶方阵，要求构造矩阵B，满足 B @ X.ravel() == Y.ravel()。
该需求是我正在开发的大型代码中的一个小模块，下方提供了最小可复现代码示例，需要补全create_B函数，确保代码中的断言不会抛出错误。我个人尝试过的失败方案太多，就不在问题中罗列了。

import numpy as np


def create_B(A):
    B = np.zeros((A.size, A.size))
    return B


A = np.random.uniform(-1, 1, (100, 100))
X = np.random.uniform(-1, 1, (100, 100))

Y = A.T @ X + X @ A

B = create_B(A)

assert np.allclose(B @ X.ravel(), Y.ravel())

需要注意的是，我可以通过Python循环轻松实现该功能，但实际使用场景中A的维度非常大，且需要多次执行该操作，因此需要无Python循环的优化实现。
我理想中的解决方案形式如下，核心是构造出索引元组indices1和indices2：

def create_B(A):
    B = np.zeros((A.size, A.size))
    B[indices1] = A
    B[indices2] += A
    return B

下方给出循环版本的参考实现，大家可以自行验证其正确性：

def create_B(A):
    B = np.zeros((A.size, A.size))
    indices = np.arange(A.size).reshape(A.shape)
    for i, k in enumerate(indices):
        for j, k in enumerate(k):
            B[k, indices[:, j]] = A[:, i]
            B[k, indices[i, :]] += A[:, j]
    return B

这个问题本质是矩阵向量化的经典克罗内克积应用场景，提供两种完全无Python循环的实现：

方案1：克罗内克积实现（性能最优） #

矩阵向量化有现成的公式：vec(A @ X @ B) = (B.T ⊗ A) vec(X)，其中⊗为克罗内克积，对应numpy的np.kron函数。
对你的等式 Y = AᵀX + XA 做向量化推导可得：
vec(Y) = (I⊗Aᵀ + Aᵀ⊗I) vec(X)，其中I为单位矩阵
直接按公式实现即可：

import numpy as np

def create_B(A):
    n = A.shape[0]
    I = np.eye(n, dtype=A.dtype)
    return np.kron(I, A.T) + np.kron(A.T, I)

该实现底层完全调用numpy优化的C级运算，性能远高于Python循环版本，适合大矩阵高频调用场景。

方案2：索引构造实现（符合你要求的形式） #

如果不想依赖克罗内克积接口，可以通过numpy广播生成全局索引，完全无循环实现：

def create_B(A):
    n = A.shape[0]
    B = np.zeros((n*n, n*n), dtype=A.dtype)
    i, j = np.indices((n, n))
    rows = (i * n + j).ravel()
    # 对应A.T@X的项赋值
    cols1 = (j.reshape(-1, 1) * n + i).ravel()
    B[rows.repeat(n), cols1] = A.T.ravel()
    # 对应X@A的项累加
    cols2 = (i.reshape(-1, 1) * n + j).ravel()
    B[rows.repeat(n), cols2] += A.ravel()
    return B

两种方案均可以通过你给出的断言测试。

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