如何为numpy对象添加自定义方法以简化kronecker乘积的调用写法

你要实现的a.k(b)式调用本质是给numpy的内置类新增自定义方法，有三种常见实现方案：

方案1：全局动态绑定（猴子补丁，最便捷） #

直接在运行时给numpy的矩阵/数组类绑定自定义方法，修改全局生效：

import numpy as np

# 定义自定义k方法
def kron_method(self, other):
    return np.kron(self, other)

# 绑定到np.matrix类（适配你当前用的矩阵类型）
np.matrix.k = kron_method
# 如果需要适配普通numpy数组，也可以绑定到ndarray类
np.ndarray.k = kron_method

# 测试效果
a = np.matrix('1 0; 0 1')
b = np.matrix('1 0; 0 1')
C = a.k(b)

注意：该修改仅在当前Python进程生效，重启后会重置。如果需要长期使用可以把这段绑定代码放到你项目的初始化脚本中。

方案2：自定义子类继承（无全局污染） #

如果不想修改numpy原生类的行为，可以自己封装继承自numpy矩阵的子类：

import numpy as np

class CustomMatrix(np.matrix):
    def k(self, other):
        # 返回值保持CustomMatrix类型，支持链式调用
        return CustomMatrix(np.kron(self, other))

# 测试效果
a = CustomMatrix('1 0; 0 1')
b = CustomMatrix('1 0; 0 1')
C = a.k(b)

注意：该方案仅你自己实例化的CustomMatrix对象会带有k方法，不会影响其他地方原生numpy矩阵的行为，适合多人协作的项目使用。

方案3：可插拔补丁封装 #

如果需要灵活控制方法的启用/禁用，可以封装成补丁函数：

import numpy as np

def add_kron_method():
    if hasattr(np.matrix, 'k'):
        return
    def kron_method(self, other):
        return np.kron(self, other)
    np.matrix.k = kron_method
    np.ndarray.k = kron_method

def remove_kron_method():
    if hasattr(np.matrix, 'k'):
        del np.matrix.k
    if hasattr(np.ndarray, 'k'):
        del np.ndarray.k

# 启用方法
add_kron_method()
# 不需要时可以调用remove_kron_method()移除

注意事项 #

• 如果你需要返回值保持固定类型，可以在自定义方法中对np.kron的返回值做类型转换，比如包一层np.matrix()
• 全局猴子补丁不建议在公共开源库、生产环境核心代码中使用，容易和其他第三方库的补丁产生命名冲突，个人脚本使用完全没问题

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