jax.numpy.arctan2

jax.numpy.arctan2(x1, x2, /)

计算 x1/x2 的反正切，并选择正确的象限。

JAX 对 numpy.arctan2() 的实现

参数:

• x1 (ArrayLike) – 分子数组。

• x2 (ArrayLike) – 分母数组；应与 x1 兼容进行广播。

返回:

x1 / x2 的元素级反正切，跟踪正确的象限。

返回类型:

Array

另请参阅

• jax.numpy.tan(): 计算角度的正切

• jax.numpy.atan2(): 此函数的数组 API 版本。

示例

考虑一个介于 0 和 \(2\pi\) 之间的弧度角序列

>>> theta = jnp.linspace(-jnp.pi, jnp.pi, 9)
>>> with jnp.printoptions(precision=2, suppress=True):
...   print(theta)
[-3.14 -2.36 -1.57 -0.79  0.    0.79  1.57  2.36  3.14]

这些角度可以等效地表示为单位圆上的 (x, y) 坐标

>>> x, y = jnp.cos(theta), jnp.sin(theta)

为了重建输入角度，我们可能会倾向于使用恒等式 \(\tan(\theta) = y / x\)，并计算 \(\theta = \tan^{-1}(y/x)\)。不幸的是，这并不能恢复输入角度

>>> with jnp.printoptions(precision=2, suppress=True):  
...   print(jnp.arctan(y / x))
[-0.    0.79  1.57 -0.79  0.    0.79  1.57 -0.79  0.  ]

问题在于 \(y/x\) 存在一些歧义：尽管 \((y, x) = (-1, -1)\) 和 \((y, x) = (1, 1)\) 在笛卡尔空间中代表不同的点，但在两种情况下 \(y / x = 1\)，因此简单的反正切方法会丢失关于角度所在象限的信息。arctan2() 就是为了解决这个问题而设计的

>>> with jnp.printoptions(precision=2, suppress=True):
...  print(jnp.arctan2(y, x))
[ 3.14 -2.36 -1.57 -0.79  0.    0.79  1.57  2.36 -3.14]

结果与输入 theta 匹配，除了端点处 \(+\pi\) 和 \(-\pi\) 在单位圆上代表无法区分的点。按照约定，arctan2() 始终返回介于 \(-\pi\) 和 \(+\pi\)（含）之间的值。