功能算法

在软件中实现数学函数是一个非平凡的任务，特别是当函数必须在所有可能的输入（包括复数无穷大、极小的或极大的浮点数值）上返回正确（或非常接近正确）的结果时。此类算法通常根据输入在实线或复平面中的位置使用函数的不同近似。

本项目提供了一种定义数学函数的功能算法的工具，并为各种编程语言和数学库生成算法的实现。目标是提供在整个复平面或实线上保证产生正确值的算法。

本项目的动机源于实现各种数学库中复杂算法的需要，否则使用LISP类语言手动执行可能会是一项繁琐且容易出错的任务。例如，为Python或NumPy目标定义反正弦算法的LOC约为45，但对于StableHLO目标，LOC为186。手动实现StableHLO的正弦算法将非常困难。

支持的算法

目前，提供了以下数学函数算法定义：

• acos(z: complex | float)，使用修改后的Hull等算法用于复数acos，
• acosh(z: complex | float)，使用关系acosh(z) = sign(z.imag) * I * acos(z)，
• asin(z: complex | float)，使用改进的Hull等人算法处理复数的asin，
• asinh(z: complex | float)，使用关系asinh(z) = -I * asin(z * I)，
• atan(z: complex | float)，使用关系atan(z) = -I * atanh(z * I)，
• atanh(z: complex | float)，使用自定义算法，
• hypot(x: float, y: float)和absolute(z: complex)，
• square(z: complex | float).

支持的目标

目前，以下目标库和语言提供了支持算法的实现

• Python，使用math函数处理实数输入，
• NumPy，使用NumPy函数处理实数输入，
• StableHLO，使用现有分解和操作，
• XLA客户端，使用现有分解和操作，
• C++，使用C++标准库。

Python和NumPy目标仅用于调试和测试功能算法。

结果

查看结果，以获取支持目标上所有支持算法的生成实现。请随意将生成的源代码复制到您的程序中，前提是与本项目相同的许可条件。

测试算法和生成实现

为确保提供的算法的正确性和准确性，我们将使用MPMath作为数学函数的参考库。我们假设mpmath实现产生正确的结果，用于任意精度的数学函数 - 这是确保准确性的先决条件。为确保正确性，我们将针对每个函数情况分别验证此假设，以消除由于MPMath中可能存在的错误导致的假阳性。

通常使用32位和64位浮点数及其复数组合通过NumPy验证算法。对NumPy目标实现的评估是在对数均匀样本上进行的（相邻样本的ULP距离是恒定的），这些样本代表整个复平面或实线，包括复数无穷大，以及极小和极大值。

为了表征算法的正确性，我们将使用ULP来衡量函数结果与其参考值之间的距离。

当使用超过一百万个在复平面或实线上对数均匀分布的样本时，以下表格显示了支持算法的函数返回值与参考值相差给定数量ULP的概率

查看所有提供算法的ULP差异的完整表格。

案例研究：square

平方函数的朴素实现可以定义为

def square(z):
    return z * z

这在实线上产生正确的结果，然而，对于复数输入，在复平面上存在区域，使用普通复数乘法的给定算法会产生错误的值。例如

>>> def square(z):
...   return z * z
... 
>>> z = complex(1e170, 1e170)
>>> square(z)
(nan+infj)

其中由于1e170 * 1e170的溢出而期望虚部为inf，但实部应该是零，但这里的实部nan来自以下计算1e170 * 1e170 - 1e170 * 1e170 -> inf - inf -> nan。

顺便说一下，我们不能依赖于NumPy的平方函数作为参考，因为它也会产生错误值（可能是平台依赖的方式）

>>> numpy.square(z)
(-inf+infj)

在这个项目中，平方函数使用以下算法

def square(ctx, z):
    if z.is_complex:
        real = ctx.select(abs(z.real) == abs(z.imag), 0, (z.real - z.imag) * (z.real + z.imag))
        imag = 2 * (z.real * z.imag)
        return ctx.complex(real, imag)
    return z * z

从这个算法可以生成不同库和编程语言的实现。例如，为了生成Python的平方函数，我们将使用

>>> import functional_algorithms as fa
>>> ctx = fa.Context()
>>> square_graph = ctx.trace(square, complex)
>>> py_square = fa.targets.python.as_function(square_graph)
>>> py_square(z)
infj

一般来说，py_square在整个复平面上产生正确的结果。

深入了解细节

让我们看看上述示例的一些细节。首先，square_graph是一个代表使用纯函数形式跟踪函数的Expr实例

>>> print(square_graph)
(def square, (z: complex),
  (complex
    (select
      (eq
        (abs (real z)),
        (abs (imag z))),
      (constant 0, (real z)),
      (multiply
        (subtract
          (real z),
          (imag z)),
        (add
          (real z),
          (imag z)))),
    (multiply
      (constant 2, (real z)),
      (multiply
        (real z),
        (imag z)))))

模块对象fa.targets.python定义了所谓的Python目标实现。存在其他目标，例如fa.targets.numpy、fa.targets.stablehlo等。

要可视化给定目标的实现，比如fa.targets.python，我们将使用tostring(<target>)方法

>>> print(square_graph.tostring(fa.targets.python))
def square(z: complex) -> complex:
    real_z: float = (z).real
    imag_z: float = (z).imag
    return complex(
        (0) if ((abs(real_z)) == (abs(imag_z))) else (((real_z) - (imag_z)) * ((real_z) + (imag_z))),
        (2) * ((real_z) * (imag_z)),
    )

这实际上是上面评估py_square(z)时使用的Python函数的定义。

同样，我们可以为其他目标生成实现，例如

>>> print(square_graph.tostring(fa.targets.stablehlo))
def : Pat<(CHLO_Square ComplexElementType:$z),
  (StableHLO_ComplexOp
    (StableHLO_SelectOp
      (StableHLO_CompareOp
       (StableHLO_AbsOp
         (StableHLO_RealOp:$real_z $z)),
       (StableHLO_AbsOp
         (StableHLO_ImagOp:$imag_z $z)),
        StableHLO_ComparisonDirectionValue<"EQ">,
        (STABLEHLO_DEFAULT_COMPARISON_TYPE)),
      (StableHLO_ConstantLike<"0"> $real_z),
      (StableHLO_MulOp
        (StableHLO_SubtractOp $real_z, $imag_z),
        (StableHLO_AddOp $real_z, $imag_z))),
    (StableHLO_MulOp
      (StableHLO_ConstantLike<"2"> $real_z),
      (StableHLO_MulOp $real_z, $imag_z)))>;

在NumPy目标的情况下，参数类型必须包括位宽信息

>>> np_square_graph = ctx.trace(square, numpy.complex64)
>>> print(np_square_graph.tostring(fa.targets.numpy))
def square(z: numpy.complex64) -> numpy.complex64:
    with warnings.catch_warnings(action="ignore"):
        z = numpy.complex64(z)
        real_z: numpy.float32 = (z).real
        imag_z: numpy.float32 = (z).imag
        result = make_complex(
            (
                (numpy.float32(0))
                if (numpy.equal(numpy.abs(real_z), numpy.abs(imag_z), dtype=numpy.bool_))
                else (((real_z) - (imag_z)) * ((real_z) + (imag_z)))
            ),
            (numpy.float32(2)) * ((real_z) * (imag_z)),
        )
        return result

>>> fa.targets.numpy.as_function(np_square_graph)(z)
infj

实用提示

调试NumPy目标实现

tostring方法中的一个有用特性是debug关键字参数。当它大于0时，将在函数实现中插入类型检查语句

>>> print(np_square_graph.tostring(fa.targets.numpy, debug=1))
def square(z: numpy.complex64) -> numpy.complex64:
    with warnings.catch_warnings(action="ignore"):
        z = numpy.complex64(z)
        real_z: numpy.float32 = (z).real
        assert real_z.dtype == numpy.float32, (real_z.dtype, numpy.float32)
        imag_z: numpy.float32 = (z).imag
        assert imag_z.dtype == numpy.float32, (imag_z.dtype, numpy.float32)
        result = make_complex(
            (
                (numpy.float32(0))
                if (numpy.equal(numpy.abs(real_z), numpy.abs(imag_z), dtype=numpy.bool_))
                else (((real_z) - (imag_z)) * ((real_z) + (imag_z)))
            ),
            (numpy.float32(2)) * ((real_z) * (imag_z)),
        )
        assert result.dtype == numpy.complex64, (result.dtype,)
        return result

当debug=2时，在调用函数时将打印出函数实现源代码和所有变量的值

>>> fa.targets.numpy.as_function(np_square_graph, debug=2)(3 + 4j)
def square(z: numpy.complex64) -> numpy.complex64:
    with warnings.catch_warnings(action="ignore"):
        z = numpy.complex64(z)
        print("z=", z)
        real_z: numpy.float32 = (z).real
        print("real_z=", real_z)
        assert real_z.dtype == numpy.float32, (real_z.dtype, numpy.float32)
        imag_z: numpy.float32 = (z).imag
        print("imag_z=", imag_z)
        assert imag_z.dtype == numpy.float32, (imag_z.dtype, numpy.float32)
        result = make_complex(
            (
                (numpy.float32(0))
                if (numpy.equal(numpy.abs(real_z), numpy.abs(imag_z), dtype=numpy.bool_))
                else (((real_z) - (imag_z)) * ((real_z) + (imag_z)))
            ),
            (numpy.float32(2)) * ((real_z) * (imag_z)),
        )
        print("result=", result)
        assert result.dtype == numpy.complex64, (result.dtype,)
        return result

z= (3+4j)
real_z= 3.0
imag_z= 4.0
result= (-7+24j)
(-7+24j)

Python和NumPy目标实现中的中间变量

在生成实现时，可以控制中间变量的命名以及它们的外观。默认情况下，仅生成用于多次作为子表达式的表达式的中继变量。但是，也可以强制为更好的实现可视化创建中继变量。为此，我们将重新定义平方算法如下

def square(ctx, z):
    if z.is_complex:
        x = abs(z.real)
        y = abs(z.imag)
        real = ctx.select(x == y, 0, ((x - y) * (y + y)).reference("real_part"))
        imag = 2 * (x * y)
        r = ctx.complex(real.reference(), imag.reference())
        return ctx(r)
    return z * z

注意使用reference方法强制将表达式定义为变量。还要注意将返回值用ctx(...)调用包装，这将赋予函数中的变量以表达式的引用值。

上述定义的Python目标的生成实现是

>>> square_graph = ctx.trace(square, complex)
>>> print(square_graph.tostring(fa.targets.python))
def square(z: complex) -> complex:
    real_z: float = (z).real
    x: float = abs(real_z)
    y: float = abs((z).imag)
    real_part: float = ((x) - (y)) * ((y) + (y))
    real: float = (0) if ((x) == (y)) else (real_part)
    imag: float = (2) * ((x) * (y))
    return complex(real, imag)

这比上面显示的实现更具表达性。