jax.numpy.correlate

jax.numpy.correlate(a, v, mode='valid', *, precision=None, preferred_element_type=None)

两个一维数组的卷积。

JAX 对 numpy.correlate() 的实现。

一维数组的卷积定义为

\[c_k = \sum_j a_{k + j} \overline{v_j}\]

其中 \(\overline{v_j}\) 是 \(v_j\) 的复共轭。

参数:

• a (ArrayLike) – 卷积的左侧输入。必须具有 a.ndim == 1。

• v (ArrayLike) – 卷积的右侧输入。必须具有 v.ndim == 1。

• mode (str) –

  控制输出的大小。可用操作包括

  • "full": 输出输入的完整卷积。

  • "same": 返回 "full" 输出的居中部分，大小与 a 相同。

  • "valid": (默认) 返回 "full" 输出中不依赖于数组边缘填充的部分。

• precision (lax.PrecisionLike) – 指定计算的精度。有关可用值的描述，请参阅 jax.lax.Precision。

• preferred_element_type (DTypeLike | None) – 数据类型，指示将结果累积到该类型并返回具有该数据类型的结果。默认值为 None，表示输入类型的默认累积类型。

返回:

包含互相关结果的数组。

返回类型:

Array

另请参阅

• jax.scipy.signal.correlate(): ND 卷积

• jax.numpy.convolve(): 一维卷积

示例

>>> x = jnp.array([1, 2, 3, 2, 1])
>>> y = jnp.array([4, 5, 6])

由于默认 mode = 'valid'，jax.numpy.correlate 只返回两个数组完全重叠的卷积部分。

>>> jnp.correlate(x, y)
Array([32., 35., 28.], dtype=float32)

指定 mode = 'full' 会使用边缘的隐式零填充来返回完整的卷积。

>>> jnp.correlate(x, y, mode='full')
Array([ 6., 17., 32., 35., 28., 13.,  4.], dtype=float32)

指定 mode = 'same' 会返回一个与第一个输入大小相同的居中卷积。

>>> jnp.correlate(x, y, mode='same')
Array([17., 32., 35., 28., 13.], dtype=float32)

如果两个输入数组都是实值且对称，则结果也将是对称的，并且等于 jax.numpy.convolve 的结果。

>>> x1 = jnp.array([1, 2, 3, 2, 1])
>>> y1 = jnp.array([4, 5, 4])
>>> jnp.correlate(x1, y1, mode='full')
Array([ 4., 13., 26., 31., 26., 13.,  4.], dtype=float32)
>>> jnp.convolve(x1, y1, mode='full')
Array([ 4., 13., 26., 31., 26., 13.,  4.], dtype=float32)

对于复数值输入

>>> x2 = jnp.array([3+1j, 2, 2-3j])
>>> y2 = jnp.array([4, 2-5j, 1])
>>> jnp.correlate(x2, y2, mode='full')
Array([ 3. +1.j,  3.+17.j, 18.+11.j, 27. +4.j,  8.-12.j], dtype=complex64)