Repository: pyfar/pyfar
Pull request #719 (2024-11-08 11:48:46)

Feature: correlate
Was merged into main by mistake in #713. Now merging into dev


Patch:
diff --git a/pyfar/dsp/__init__.py b/pyfar/dsp/__init__.py
index bc83ad0d..2acbf014 100644
--- a/pyfar/dsp/__init__.py
+++ b/pyfar/dsp/__init__.py
@@ -24,6 +24,7 @@ from .dsp import (
     rms,
     normalize,
     average,
+    correlate,
 )
 
 from .interpolation import (
@@ -67,4 +68,5 @@ __all__ = [
     'average',
     'normalize',
     'fractional_time_shift',
+    'correlate',
 ]
diff --git a/pyfar/dsp/dsp.py b/pyfar/dsp/dsp.py
index b44dbac7..fb77a106 100644
--- a/pyfar/dsp/dsp.py
+++ b/pyfar/dsp/dsp.py
@@ -2667,3 +2667,194 @@ def normalize(signal, reference_method='max', domain='auto',
         return normalized_signal, reference_norm
     else:
         return normalized_signal
+
+
+def correlate(signal_1, signal_2, mode='full', normalize=False):
+    r"""
+    Compute the channel-wise correlation function between signals.
+
+    The correlation function of the time signals :math:`x_1[n]` and
+    :math:`x_2[n]` is given by
+
+    .. math:: c[l] = \sum_n x_1[n] x_2[n-l]
+
+    with the index :math:`n` and time lag :math:`l` in samples. The computation
+    is realized in the frequency domain using the corresponding spectra
+    :math:`X_1[k]` and :math:`X_2[k]`
+
+    .. math:: c = \mathrm{IFFT\{X_1 X_2^-\}}
+
+    with the inverse fourier transform denoted by IFFT and :math:`X^-` denoting
+    the spectrum of a time-reversed and complex conjugated time signal, i.e.,
+    ``X_minus = fft(conj(x[::-1]))`` where the conjugate is required if the
+    input has complex-valued time data.
+
+    Parameters
+    ----------
+    signal_1 : Signal
+        The first input signal. It must have the same sample rate as
+        `signal_2` and its ``cshape`` must be broadcastable to that of
+        `signal_2`.
+    signal_2 : Signal
+        The second input signal. It must have the same sample rate as
+        `signal_1` and its ``cshape`` must be broadcastable to that of
+        `signal_1`.
+    mode : str, optional
+        Specifies how the correlation is computed.
+
+        ``'full'``
+            Computes the full correlation function by zero padding the input
+            signals to a length of
+            ``signal_1.n_samples + signal_2.n_samples - 1`` before applying the
+            Fourier transform (see equations above).
+        ``'cyclic'``
+            Computes the cyclic correlation function, which uses the input
+            signals as they are. In this case `signal_1` and `signal_2` must
+            have the same number of samples (length).
+
+        The default is ``'full'``.
+    normalize : bool, optional
+        If ``True``, the correlation function is normalized to force
+        :math:`|c[l]|\leq1`. This is done by the division
+        :math:`c[l]/\sqrt{E(x_1[n])\,E(x_2[n])}`, where :math:`E(\cdot)`
+        denotes the :py:func:`~energy`. The default ``False`` does not apply
+        any normalization. Normalization is not available for complex-valued
+        time signals.
+
+    Returns
+    -------
+    correlation : TimeData
+        The correlation function :math:`c[l]` is contained in
+        ``correlation.time`` and the lags in seconds, i.e., the delays applied
+        to `signal_2` (see equations above) are contained in
+        ``correlation.times``. ``correlation.time`` is complex if one of the
+        input signals has complex-valued time data. The cshape of `correlation`
+        matches the cshape to which `signal_1` and `signal_2` were broadcasted.
+        The lags can be converted to samples by multiplication with
+        ``signal_1.sampling_rate``. In this case, they are in the interval
+        ``[-signal_2.n_samples + 1, signal_1.n_samples - 1]`` if the full
+        correlation was computed. In case of the cyclic correlation, they are
+        in the interval
+        ``[-(signal_1.n_samples // 2) + 1, signal_1.n_samples // 2]``
+        if the signals have an even number of samples, and in the interval
+        ``[-(signal_1.n_samples // 2), signal_1.n_samples // 2]``
+        if the signals have an odd number of samples.
+
+    Examples
+    --------
+    Compute the lags (delay) that are required to maximize the
+    cross-correlation between a one and multi-dimensional signal
+
+    .. plot::
+
+        >>> import pyfar as pf
+        >>> import numpy as np
+        >>>
+        >>> # one-dimensional signal: impulse with zero delay
+        >>> signal_1 = pf.signals.impulse(5, 0)
+        >>>
+        >>> # multi-dimension signal with cshape = (2, 2):
+        >>> # impulses with non-zero delays
+        >>> delays = np.array([[0, 1], [2, 3]], dtype=int)
+        >>> signal_2 = pf.signals.impulse(5, delays)
+        >>>
+        >>> cor = pf.dsp.correlate(signal_1, signal_2, 'full')
+        >>>
+        >>> # compute the lags in samples
+        >>> argmax = cor.times[np.argmax(cor.time, axis=-1)]
+        >>>
+        >>> # plot correlation and indicate maximum
+        >>> ax = pf.plot.time(cor, unit='ms')
+        >>> ax.set_title('Correlation and position of maxima (dots)')
+        >>> ax.set_xlabel('Time lag in ms')
+        >>> ax.set_ylabel('Auto correlation')
+        >>> for amax, color in zip(argmax.flatten(), 'bryp'):
+        >>>     ax.axvline(amax, color=pf.plot.color(color), linestyle=':')
+
+    Linear and cyclic auto-correlation of a perfect sequence. Perfect sequences
+    have unit auto-correlation for :math:`l=0` and zero auto-correlation
+    otherwise
+
+    .. plot::
+
+        >>> import pyfar as pf
+        >>>
+        >>> signal = pf.signals.linear_perfect_sweep(2**7)
+        >>>
+        >>> for mode in ['full', 'cyclic']:
+        >>>     line = '--' if mode == 'cyclic' else '-'
+        >>>     cor = pf.dsp.correlate(signal, signal, mode, normalize=True)
+        >>>     ax = pf.plot.time(cor, unit='ms', label=mode, ls=line)
+        >>>
+        >>> ax.set_xlabel('Time lag in ms')
+        >>> ax.set_ylabel('Auto correlation')
+        >>> ax.legend()
+    """
+
+    # check input
+    if type(signal_1) is not pyfar.Signal or \
+            type(signal_2) is not pyfar.Signal:
+        raise TypeError("signal_1 and signal_2 must be pyfar.Signal objects")
+
+    if signal_1.sampling_rate != signal_2.sampling_rate:
+        raise ValueError("Both signals must have the same sampling rate.")
+
+    if mode not in ['full', 'cyclic']:
+        raise ValueError(f"mode is '{mode}' but must be 'full' or 'cyclic'")
+
+    # copy input to avoid changing mutual data
+    signal_1 = signal_1.copy()
+    signal_2 = signal_2.copy()
+
+    # check for complex signals
+    if signal_1.complex and not signal_2.complex:
+        signal_2.complex = True
+    if signal_2.complex and not signal_1.complex:
+        signal_1.complex = True
+
+    # compute normalization factor before zero padding
+    if normalize:
+        normalization = np.sqrt(energy(signal_1) * energy(signal_2))
+
+    # determine signal length for the FFT
+    n_samples = np.array([signal_1.n_samples, signal_2.n_samples], dtype=int)
+    if mode == 'cyclic':
+        if n_samples[0] != n_samples[1]:
+            raise ValueError(("signal_1 and signal_2 must be of the same "
+                              "length in 'cyclic' mode"))
+        n_fft = n_samples[0]
+    else:
+        n_fft = np.sum(n_samples) - 1
+
+        # zero pad signals to match lengths
+        if signal_1.n_samples < n_fft:
+            signal_1 = pyfar.dsp.pad_zeros(
+                signal_1, n_fft - signal_1.n_samples, 'end')
+        if signal_2.n_samples < n_fft:
+            signal_2 = pyfar.dsp.pad_zeros(
+                signal_2, n_fft - signal_2.n_samples, 'beginning')
+
+    # compute correlation as frequency domain convolution
+    # with time flipped and conjugate values for second signal
+    signal_2.time = signal_2.time[..., ::-1].conj()
+    correlation = signal_1 * signal_2
+    correlation = correlation.time
+
+    # apply normalization
+    if normalize:
+        correlation /= normalization
+
+    # compute lags, i.e., times that the second signal was shifted
+    # with respect to the first
+    if mode == 'cyclic':
+        roll = n_fft // 2
+        correlation = np.roll(correlation, -roll, -1)
+        if n_fft % 2:
+            lags = np.arange(-roll, roll + 1)
+        else:
+            lags = np.arange(-roll + 1, roll + 1)
+    else:
+        lags = np.arange(-n_samples[1] + 1, n_samples[0])
+
+    return pyfar.TimeData(correlation, lags / signal_1.sampling_rate,
+                          is_complex=signal_1.complex)
diff --git a/pyfar/signals/deterministic.py b/pyfar/signals/deterministic.py
index 6dc49dec..7a95163f 100644
--- a/pyfar/signals/deterministic.py
+++ b/pyfar/signals/deterministic.py
@@ -649,20 +649,12 @@ def linear_perfect_sweep(
         >>> sweep = pf.signals.linear_perfect_sweep(2**8)
         >>>
         >>> # compute auto-correlation
-        >>> auto_correlation = np.empty(2**8)
-        >>> for idx, shift in enumerate(range(-2**7, 2**7)):
-        >>>     auto_correlation[idx] = np.dot(
-        ...         sweep.time.flatten(),
-        ...         np.roll(sweep.time.flatten(), shift))
+        >>> auto_correlation = pf.dsp.correlate(
+        ...     sweep, sweep, 'cyclic', normalize=True)
         >>>
-        >>> auto_correlation /= pf.dsp.energy(sweep)
-        >>>
-        >>> # plot auto-correlation
-        >>> with pf.plot.context():
-        >>>     plt.plot(np.arange(-2**7, 2**7), auto_correlation)
-        >>>     plt.gca().set_xlim(-2**7, 2**7)
-        >>>     plt.gca().set_xlabel('time lag in samples')
-        >>>     plt.gca().set_ylabel('auto correlation')
+        >>> ax = pf.plot.time(auto_correlation, unit='ms')
+        >>> ax.set_xlabel('Time lag in ms')
+        >>> ax.set_ylabel('Auto correlation')
     """
 
     return _frequency_domain_sweep(


Test patch:
diff --git a/tests/test_dsp_correlate.py b/tests/test_dsp_correlate.py
new file mode 100644
index 00000000..8fa68369
--- /dev/null
+++ b/tests/test_dsp_correlate.py
@@ -0,0 +1,206 @@
+from pyfar.dsp import correlate
+import pyfar as pf
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+
+def test_error_signal_type():
+    """Test if a TypeError is raised for non Signal input."""
+
+    signal = pf.Signal(1, 1)
+
+    with pytest.raises(TypeError, match="signal_1 and signal_2 must be"):
+        correlate('not a signal', signal)
+
+    with pytest.raises(TypeError, match="signal_1 and signal_2 must be"):
+        correlate(signal, 'not a signal')
+
+    with pytest.raises(TypeError, match="signal_1 and signal_2 must be"):
+        correlate('not a signal', 'not a signal')
+
+
+def test_error_different_sampling_rates():
+    """Test if signals with different sampling rates raise an error."""
+
+    signal_1 = pf.Signal(1, 1)
+    signal_2 = pf.Signal(1, 2)
+
+    with pytest.raises(ValueError, match="same sampling rate"):
+        correlate(signal_1, signal_2)
+
+
+def test_error_mode_string():
+    """Test if a ValueError is raised for invalid mode values."""
+
+    signal = pf.Signal(1, 1)
+
+    with pytest.raises(ValueError, match="invalid_mode"):
+        correlate(signal, signal, 'invalid_mode')
+
+
+def test_error_cyclic_mode():
+    """
+    Test if ValueError for signals with differing length is raised in cyclic
+    mode.
+    """
+
+    signal_1 = pf.Signal(1, 1)
+    signal_2 = pf.Signal([1, 2], 1)
+
+    with pytest.raises(ValueError, match="must be of the same length"):
+        correlate(signal_1, signal_2, mode='cyclic')
+
+
+def test_mutable_objects():
+    """Test if input does not change."""
+
+    signal_1 = pf.Signal([1, 0, 0], 1)
+    signal_2 = pf.Signal([2, 0, 0], 1)
+
+    correlate(signal_1, signal_2)
+
+    assert signal_1 == pf.Signal([1, 0, 0], 1)
+    assert signal_2 == pf.Signal([2, 0, 0], 1)
+
+
+@pytest.mark.parametrize('length_1', [4, 5])
+@pytest.mark.parametrize('delay_1', [0, 1, 2])
+@pytest.mark.parametrize('length_2', [4, 5])
+@pytest.mark.parametrize('delay_2', [0, 1, 2])
+def test_full_mode_1d(length_1, delay_1, length_2, delay_2):
+    """Test full mode for different combinations of 1D signals."""
+
+    # compute correlation
+    signal_1 = pf.signals.impulse(length_1, delay_1, sampling_rate=1)
+    signal_2 = pf.signals.impulse(length_2, delay_2, sampling_rate=1)
+
+    correlation = correlate(signal_1, signal_2)
+
+    # test output format
+    assert type(correlation) == pf.TimeData
+    assert correlation.time.shape == (1, length_1 + length_2 - 1)
+    assert correlation.times.shape == (length_1 + length_2 - 1, )
+
+    # test lag that maximizes the correlation
+    lag = (delay_1 - delay_2)
+
+    # test correlation and lags
+    npt.assert_almost_equal(
+        correlation.time[0, correlation.times==lag], 1., 10)
+    npt.assert_almost_equal(
+        correlation.time[0, correlation.times!=lag], 0., 10)
+
+
+@pytest.mark.parametrize('length', [4, 5])
+@pytest.mark.parametrize('delay_1', [0, 1, 2])
+@pytest.mark.parametrize('delay_2', [0, 1, 2])
+def test_cyclic_mode_1d(length, delay_1, delay_2):
+    """Test cyclic mode for different combinations of 1D signals."""
+
+    # compute correlation
+    signal_1 = pf.signals.impulse(length, delay_1, sampling_rate=1)
+    signal_2 = pf.signals.impulse(length, delay_2, sampling_rate=1)
+
+    correlation = correlate(signal_1, signal_2, mode='cyclic')
+
+    # test output format
+    assert type(correlation) == pf.TimeData
+    assert correlation.time.shape == (1, length)
+    assert correlation.times.shape == (length, )
+
+    # test lag that maximizes the correlation
+    # (lag must be adjusted to cyclic case)
+    lag = delay_1 - delay_2
+    if length % 2 and lag < -(length//2):
+        lag = length + lag
+    if not length % 2 and lag <= -(length//2):
+        lag = length + lag
+
+    # test correlation and lags
+    npt.assert_almost_equal(
+        correlation.time[0, correlation.times==lag], 1., 10)
+    npt.assert_almost_equal(
+        correlation.time[0, correlation.times!=lag], 0., 10)
+
+
+def test_full_mode_nd():
+    """Test broadcasting and n-dimensional signals in full mode."""
+
+    # compute correlation
+    delays = np.array([[0, 1], [2, 3]], dtype=int)
+    signal_1 = pf.signals.impulse(5, 0, sampling_rate=1)
+    signal_2 = pf.signals.impulse(5, delays, sampling_rate=1)
+
+    correlation = correlate(signal_1, signal_2)
+
+    # test output shapes
+    assert type(correlation) == pf.TimeData
+    assert correlation.time.shape == (2, 2, 9)
+    assert correlation.times.shape == (9, )
+
+    # test correlation and lags
+    for dim_1 in range(2):
+        for dim_2 in range(2):
+            npt.assert_almost_equal(
+                correlation.time[dim_1, dim_2,
+                                 correlation.times==-delays[dim_1, dim_2]],
+                1., 10)
+            npt.assert_almost_equal(
+                correlation.time[dim_1, dim_2,
+                                 correlation.times!=-delays[dim_1, dim_2]],
+                0., 10)
+
+
+@pytest.mark.parametrize('normalize', [False, True])
+def test_normalize(normalize):
+    """Test the normalize parameter."""
+    # sine with 4 samples period -> auto-correlation is 4 samples period sine
+    signal = pf.signals.sine(5, 8, 1, 0, 20)
+    correlation = correlate(signal, signal, 'cyclic', normalize)
+    max_correlation = 1 if normalize else pf.dsp.energy(signal)
+
+    npt.assert_almost_equal(correlation.time[0, 0::2], np.zeros(4), 10)
+    npt.assert_almost_equal(correlation.time[0, 1::4],
+                            -max_correlation * np.ones(2), 10)
+    npt.assert_almost_equal(correlation.time[0, 3::4],
+                            max_correlation * np.ones(2), 10)
+
+
+def test_complex_signals():
+    """Test with complex signals."""
+
+    signal_1 = pf.Signal([0, 0, 1+0j], 1, is_complex=True)
+    signal_2 = pf.Signal([0+1j, 0, 0], 1, is_complex=True)
+
+    correlation = correlate(signal_1, signal_2)
+
+    # test correlation values
+    lag = 2
+    npt.assert_almost_equal(
+        correlation.time[0, correlation.times==lag], 0-1j, 10)
+    npt.assert_almost_equal(
+        correlation.time[0, correlation.times!=lag], 0+0j, 10)
+
+
+@pytest.mark.parametrize('complex_first', [True, False])
+def test_complex_and_real_signals(complex_first):
+    """Test with real and complex-valued signals."""
+
+    if complex_first:
+        signal_1 = pf.Signal([0, 0, 1j], 1, is_complex=True)
+        signal_2 = pf.Signal([1, 0, 0], 1, is_complex=True)
+        expected = 1j
+    else:
+        signal_1 = pf.Signal([0, 0, 1], 1, is_complex=True)
+        signal_2 = pf.Signal([1j, 0, 0], 1, is_complex=True)
+        expected = -1j
+
+    correlation = correlate(signal_1, signal_2)
+
+    # test correlation values
+    lag = 2
+    npt.assert_almost_equal(
+        correlation.time[0, correlation.times==lag], expected, 10)
+    npt.assert_almost_equal(
+        correlation.time[0, correlation.times!=lag], 0+0j, 10)