Merge pull request #8128 from takagi/add-freq_shift

Add `cupyx.signal.freq_shift`

cupyx/signal/__init__.py

@@ -5,3 +5,4 @@
 from cupyx.signal._convolution import convolve1d3o  # NOQA
 from cupyx.signal._radartools import pulse_compression, pulse_doppler, cfar_alpha  # NOQA
 from cupyx.signal._filtering import firfilter, firfilter2, firfilter_zi  # NOQA
+from cupyx.signal._filtering import freq_shift  # NOQA

cupyx/signal/_filtering/__init__.py

@@ -1 +1,2 @@
 from cupyx.signal._filtering._filtering import firfilter, firfilter2, firfilter_zi  # NOQA
+from cupyx.signal._filtering._filtering import freq_shift  # NOQA

cupyx/signal/_filtering/_filtering.py

@@ -1,3 +1,5 @@
+import math
+
 import cupy
 import cupyx.scipy.linalg
 import cupyx.scipy.signal._signaltools
@@ -246,3 +248,30 @@ def firfilter2(
         y = axis_slice(y, start=edge, stop=-edge, axis=axis)
 
     return cupy.copy(y)
+
+
+_freq_shift_kernel = cupy.ElementwiseKernel(
+    "float64 x, float64 c", "complex128 out",
+    "out = thrust::polar(x, c * i);",
+    "_freq_shift_kernel",
+)
+
+
+def freq_shift(x, freq, fs):
+    """
+    Frequency shift signal by freq at fs sample rate
+
+    Parameters
+    ----------
+    x : array_like, complex valued
+        The data to be shifted.
+    freq : float
+        Shift by this many (Hz)
+    fs : float
+        Sampling rate of the signal
+    domain : string
+        freq or time
+    """
+    x = cupy.asarray(x)
+    c = -2 * math.pi * freq / fs
+    return _freq_shift_kernel(x, c)

docs/source/reference/ext.rst

@@ -27,7 +27,8 @@ through the courtesy of Nvidia cuSignal team.
 
    cupyx.signal.convolve1d3o
    cupyx.signal.pulse_compression
-
+   cupyx.signal.freq_shift
+
 Profiling utilities
 -------------------
 

tests/cupyx_tests/signal_tests/filtering_tests/test_filtering.py

@@ -51,3 +51,31 @@ def test_firfilter2(dtype, num_samps, filter_len, padtype):
     cpu_output = scipy.signal.filtfilt(cpu_filter, 1, cpu_sig, padtype=padtype)
     gpu_output = cupyx.signal.firfilter2(gpu_filter, gpu_sig, padtype=padtype)
     testing.assert_allclose(gpu_output, cpu_output, atol=1e-3, rtol=1e-3)
+
+
+def freq_shift_cpu(x, freq, fs):
+    """
+    Frequency shift signal by freq at fs sample rate
+    Parameters
+    ----------
+    x : array_like, complex valued
+        The data to be shifted.
+    freq : float
+        Shift by this many (Hz)
+    fs : float
+        Sampling rate of the signal
+    domain : string
+        freq or time
+    """
+    x = numpy.asarray(x)
+    return x * numpy.exp(-1j * 2 * numpy.pi * freq / fs * numpy.arange(x.size))
+
+
+@pytest.mark.parametrize("dtype", [cupy.float64, cupy.complex128])
+@pytest.mark.parametrize("num_samps", [2**8])
+@pytest.mark.parametrize("freq", numpy.fft.fftfreq(10, 0.1))
+@pytest.mark.parametrize("fs", [0.3])
+def test_freq_shift(dtype, num_samps, freq, fs):
+    cpu_output = freq_shift_cpu(freq, fs, num_samps)
+    gpu_output = cupyx.signal.freq_shift(freq, fs, num_samps)
+    testing.assert_allclose(gpu_output, cpu_output)