jax.numpy.arctan2

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

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

numpy.arctan2() 的 JAX 实现

参数:

• x1 (ArrayLike) – 分子数组。

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

返回:

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

返回类型:

数组

参见

• 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\)（包括端点）之间的值。