Numpy Enum vs String 作为函数参数

阅读更多：Numpy 教程

前言

在使用Numpy库时，经常需要在函数中传入参数，并且这个参数是一个状态、类型等取值范围有限的变量。比如，当我们需要进行数组的计算时，我们经常需要传入参数选择计算方式，这个参数就只能取几个固定值。那么，这个参数我们是使用Numpy枚举(Enum)类型还是普通字符串(String)类型呢？这个问题需要我们深入探讨。

Numpy Enum

Numpy枚举是新版本Numpy从1.20.0版本开始提供的一种数据类型。在使用时，需要从Numpy导入Enum类型。

下面是Numpy枚举的一个示例：

import numpy as np

class Operation(np.enum.IntEnum):
    Add = 0
    Sub = 1
    Mul = 2
    Div = 3

def calculate(arr: np.ndarray, op: Operation):
    if op == Operation.Add:
        return np.sum(arr)
    elif op == Operation.Sub:
        return np.subtract.reduce(arr)
    elif op == Operation.Mul:
        return np.prod(arr)
    elif op == Operation.Div:
        return np.divide.reduce(arr)
    else:
        raise ValueError("Invalid operation")

在上面的代码中，我们定义了一个枚举类型Operation，它继承自Numpy的IntEnum类型。每一个枚举项都有一个int值与之对应。在calculate函数中，我们传入一个Operation类型的参数，根据不同的取值来进行不同的计算。在代码中，我们使用了枚举变量来比较，这样可以保证类型的安全性，避免发生因字符串拼写错误导致的计算错误。

字符串作为参数

我们也可以使用普通字符串作为参数。比如，在我们上面的枚举示例中，我们可以将Operation枚举替换成字符串：

def calculate(arr: np.ndarray, op: str):
    if op == 'Add':
        return np.sum(arr)
    elif op == 'Sub':
        return np.subtract.reduce(arr)
    elif op == 'Mul':
        return np.prod(arr)
    elif op == 'Div':
        return np.divide.reduce(arr)
    else:
        raise ValueError("Invalid operation")

在上面的代码中，我们直接使用字符串来替换了枚举，其他代码与前面示例类似。在使用时，只需要传入一个字符串即可。

枚举与字符串的比较

在上面的两个示例中，我们分别使用Numpy枚举和字符串两种方式来进行计算。那么，这两种方式在性能上有什么区别呢？

为了测试这个问题，我们分别实现了两个函数，其实现方式与上面两个示例类似。

def calculate_with_enum(arr: np.ndarray, op: Operation):
    if op == Operation.Add:
        return np.sum(arr)
    elif op == Operation.Sub:
        return np.subtract.reduce(arr)
    elif op == Operation.Mul:
        return np.prod(arr)
    elif op == Operation.Div:
        return np.divide.reduce(arr)
    else:
        raise ValueError("Invalid operation")

def calculate_with_str(arr: np.ndarray, op: str):
    if op == 'Add':
        return np.sum(arr)
    elif op == 'Sub':
        return np.subtract.reduce(arr)
    elif op == 'Mul':
        return np.prod(arr)
    elif op == 'Div':
        return np.divide.reduce(arr)
    else:
        raise ValueError("Invalid operation")

我们使用一个长度为10000的数组来进行计算，并对这两个函数进行性能测试。

在测试时，我们使用Jupyter Notebook中的%timeit命令来计时。测试结果如下：

arr = np.random.rand(10000)

%timeit calculate_with_enum(arr, Operation.Add)
%timeit calculate_with_str(arr, 'Add')

%timeit calculate_with_enum(arr, Operation.Sub)
%timeit calculate_with_str(arr, 'Sub')

%timeit calculate_with_enum(arr, Operation.Mul)
%timeit calculate_with_str(arr, 'Mul')