如何让自定义Python Quantity类兼容Numpy数学函数？

定义的Quantity类 #

class Quantity:
    def __init__(self, value, units):
        self.value = value
        self.units = units

需求目标 #

希望实现两种场景的自动转换：

• 传入单个Quantity实例时：np.[数学函数](a) 等价于 np.[数学函数](a.value)
• 传入Quantity实例列表时：np.[数学函数]([a1, a2]) 等价于 np.[数学函数]([a1.value, a2.value])

尝试的方案及问题 #

通过__array_ufunc__实现，但仅支持单个实例，传入列表时报错：

尝试的代码

def __array_ufunc__(self, ufunc, method, *args, **kwargs):
    out = []
    for a in args:
        if isinstance(a, self.__class__):
            out.append(a.value)
        else:
            out.append(a)
    return ufunc(out, kwargs)    

报错信息

TypeError: loop of ufunc does not support argument 0 of type Quantity which has no callable [数学函数] method

逐个定义类方法（如def sqrt(self): return np.sqrt(self.value)）可支持列表，但需重复编写代码，过于繁琐。

通过完善__array_ufunc__并补充__array__方法，可实现无需逐个定义函数的通用支持：

1. 完整实现代码 #

import numpy as np

class Quantity:
    def __init__(self, value, units):
        self.value = value
        self.units = units

    def __array__(self, dtype=None):
        # 让NumPy能直接提取实例的value作为数组元素
        return np.asarray(self.value, dtype=dtype)

    def __array_ufunc__(self, ufunc, method, *args, **kwargs):
        # 递归处理所有参数，转换Quantity实例和嵌套容器
        def convert_arg(arg):
            if isinstance(arg, Quantity):
                return arg.value
            elif isinstance(arg, (list, tuple)):
                return type(arg)(convert_arg(e) for e in arg)
            return arg

        # 转换位置参数和关键字参数
        converted_args = tuple(convert_arg(a) for a in args)
        converted_kwargs = {k: convert_arg(v) for k, v in kwargs.items()}
        
        # 调用原ufunc并返回结果
        return ufunc(*converted_args, method=method, **converted_kwargs)

2. 验证效果 #

a = Quantity(0.5, "rad")
print(np.sin(a))          # 输出: 0.479425538604203
print(np.sin([a, a]))     # 输出: [0.47942554 0.47942554]
print(np.sqrt(Quantity(4, "m")))  # 输出: 2.0

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