Develop complex signals

pyfar/dsp/dsp.py

@@ -5,6 +5,7 @@
 from pyfar.dsp import fft
 from pyfar.classes.warnings import PyfarDeprecationWarning
 import warnings
+import scipy.fft as sfft
 
 
 def phase(signal, deg=False, unwrap=False):
@@ -102,9 +103,14 @@ def group_delay(signal, frequencies=None, method='fft'):
         group_delay = group_delay.reshape(signal.cshape + (-1, ))
 
     elif method == 'fft':
-        freq_k = fft.rfft(signal.time * np.arange(signal.n_samples),
-                          signal.n_samples, signal.sampling_rate,
-                          fft_norm='none')
+        if signal.complex:
+            freq_k = fft.fft(signal.time * np.arange(signal.n_samples),
+                             signal.n_samples, signal.sampling_rate,
+                             fft_norm='none')
+        else:
+            freq_k = fft.rfft(signal.time * np.arange(signal.n_samples),
+                              signal.n_samples, signal.sampling_rate,
+                              fft_norm='none')
 
         group_delay = np.real(freq_k / signal.freq_raw)
 
@@ -276,7 +282,9 @@ def spectrogram(signal, window='hann', window_length=1024,
     Returns
     -------
     frequencies : numpy array
-        Frequencies in Hz at which the magnitude spectrum was computed
+        Frequencies in Hz at which the magnitude spectrum was computed. For
+        complex valued signals, frequencies and spectrogram is arranged such
+        that 0 Hz bin is centered.
     times : numpy array
         Times in seconds at which the magnitude spectrum was computed
     spectrogram : numpy array
@@ -298,16 +306,22 @@ def spectrogram(signal, window='hann', window_length=1024,
 
     frequencies, times, spectrogram = sgn.spectrogram(
         x=signal.time.squeeze(), fs=signal.sampling_rate, window=window,
-        noverlap=window_overlap, mode='magnitude', scaling='spectrum')
+        noverlap=window_overlap, mode='magnitude', scaling='spectrum',
+        return_onesided=not signal.complex)
 
     # remove normalization from scipy.signal.spectrogram
     spectrogram /= np.sqrt(1 / window.sum()**2)
 
     # apply normalization from signal
     if normalize:
         spectrogram = fft.normalization(
-            spectrogram, window_length, signal.sampling_rate,
-            signal.fft_norm, window=window)
+                spectrogram, window_length, signal.sampling_rate,
+                signal.fft_norm, window=window)
+
+    # rearrange spectrogram and frequencies to center 0 Hz bin
+    if signal.complex:
+        frequencies = sfft.fftshift(frequencies)
+        spectrogram = sfft.fftshift(spectrogram, axes=0)
 
     # scipy.signal takes the center of the DFT blocks as time stamp we take the
     # beginning (looks nicer in plots, both conventions are used)
@@ -1171,7 +1185,8 @@ def time_shift(
 
     if np.any(np.isnan(shifted.time)):
         shifted = pyfar.TimeData(
-            shifted.time, shifted.times, comment=shifted.comment)
+            shifted.time, shifted.times, comment=shifted.comment,
+            is_complex=signal.complex)
 
     return shifted
 
@@ -1190,7 +1205,9 @@ def find_impulse_response_delay(impulse_response, N=1):
     2.  By default a first order polynomial is used, as the slope of the
         analytic signal should in theory be linear.
 
-    Alternatively see :py:func:`pyfar.dsp.find_impulse_response_start`.
+    Alternatively see :py:func:`pyfar.dsp.find_impulse_response_start`. For
+    complex-valued time signals, the delay is calculated separately for the
+    real and complex part, and its minimum value returned.
 
     Parameters
     ----------
@@ -1238,54 +1255,65 @@ def find_impulse_response_delay(impulse_response, N=1):
     n = int(np.ceil((N+2)/2))
 
     start_samples = np.zeros(impulse_response.cshape)
+    modes = ['real', 'complex'] if impulse_response.complex else ['real']
+    start_sample = np.zeros((len(modes), 1), dtype=float)
+
     for ch in np.ndindex(impulse_response.cshape):
         # Calculate the correlation between the impulse response and its
         # minimum phase equivalent. This requires a minimum phase equivalent
         # in the strict sense, instead of the appriximation implemented in
         # pyfar.
         n_samples = impulse_response.n_samples
+        for idx, mode in enumerate(modes):
+            ir = impulse_response.time[ch]
+            ir = np.real(ir) if mode == 'real' else np.imag(ir)
+
+            if np.max(ir) > 1e-16:
+                # minimum phase warns if the input signal is not symmetric,
+                # which is not critical for this application
+                with warnings.catch_warnings():
+                    warnings.filterwarnings(
+                        "ignore", message="h does not appear to by symmetric",
+                        category=RuntimeWarning)
+                    ir_minphase = sgn.minimum_phase(
+                        ir, n_fft=4*n_samples)
+
+                correlation = sgn.correlate(
+                    ir,
+                    np.pad(ir_minphase, (0, n_samples - (n_samples + 1)//2)),
+                    mode='full')
+                lags = np.arange(-n_samples + 1, n_samples)
+
+                # calculate the analytic signal of the correlation function
+                correlation_analytic = sgn.hilbert(correlation)
+
+                # find the maximum of the analytic part of the correlation
+                # function and define the search range around the maximum
+                argmax = np.argmax(np.abs(correlation_analytic))
+                search_region_range = np.arange(argmax-n, argmax+n)
+                search_region = np.imag(
+                    correlation_analytic[search_region_range])
+
+                # mask values with a negative gradient
+                mask = np.gradient(search_region, search_region_range) > 0
+
+                # fit a polygon and estimate its roots
+                search_region_poly = np.polyfit(
+                    search_region_range[mask]-argmax, search_region[mask], N)
+                roots = np.roots(search_region_poly)
+
+                # Use only real-valued roots
+                if np.all(np.isreal(roots)):
+                    root = roots[np.abs(roots) == np.min(np.abs(roots))]
+                    start_sample[idx] = np.squeeze(lags[argmax] + root)
+                else:
+                    start_sample[idx] = np.nan
+                    warnings.warn('Starting sample not found for channel '
+                                  f'{ch}')
+            else:
+                start_sample[idx] = np.nan
 
-        # minimum phase warns if the input signal is not symmetric, which is
-        # not critical for this application
-        with warnings.catch_warnings():
-            warnings.filterwarnings(
-                "ignore", message="h does not appear to by symmetric",
-                category=RuntimeWarning)
-            ir_minphase = sgn.minimum_phase(
-                impulse_response.time[ch], n_fft=4*n_samples)
-
-        correlation = sgn.correlate(
-            impulse_response.time[ch],
-            np.pad(ir_minphase, (0, n_samples - (n_samples + 1)//2)),
-            mode='full')
-        lags = np.arange(-n_samples + 1, n_samples)
-
-        # calculate the analytic signal of the correlation function
-        correlation_analytic = sgn.hilbert(correlation)
-
-        # find the maximum of the analytic part of the correlation function
-        # and define the search range around the maximum
-        argmax = np.argmax(np.abs(correlation_analytic))
-        search_region_range = np.arange(argmax-n, argmax+n)
-        search_region = np.imag(correlation_analytic[search_region_range])
-
-        # mask values with a negative gradient
-        mask = np.gradient(search_region, search_region_range) > 0
-
-        # fit a polygon and estimate its roots
-        search_region_poly = np.polyfit(
-            search_region_range[mask]-argmax, search_region[mask], N)
-        roots = np.roots(search_region_poly)
-
-        # Use only real-valued roots
-        if np.all(np.isreal(roots)):
-            root = roots[np.abs(roots) == np.min(np.abs(roots))]
-            start_sample = np.squeeze(lags[argmax] + root)
-        else:
-            start_sample = np.nan
-            warnings.warn(f"Starting sample not found for channel {ch}")
-
-        start_samples[ch] = start_sample
+        start_samples[ch] = np.nanmin(start_sample)
 
     return start_samples
 
@@ -1656,8 +1684,11 @@ def convolve(signal1, signal2, mode='full', method='overlap_add'):
         res[..., :n_min-1] += res[..., -n_min+1:]
         res = res[..., :n_max]
 
+    is_result_complex = True if res.dtype.kind == 'c' else False
+
     return pyfar.Signal(
-        res, signal1.sampling_rate, domain='time', fft_norm=fft_norm)
+        res, signal1.sampling_rate, domain='time', fft_norm=fft_norm,
+        is_complex=is_result_complex)
 
 
 def decibel(signal, domain='freq', log_prefix=None, log_reference=1,
@@ -1814,8 +1845,14 @@ def energy(signal):
     if not isinstance(signal, pyfar.Signal):
         raise ValueError(f"signal is type '{signal.__class__}'"
                          " but must be of type 'Signal'.")
+
+    if isinstance(signal, (pyfar.Signal, pyfar.TimeData)):
+        if signal.complex:
+            raise ValueError((
+                "'energy' is not implemented for complex time signals."))
+
     # return and compute data
-    return np.sum(signal.time**2, axis=-1)
+    return np.sum(np.abs(signal.time)**2, axis=-1)
 
 
 def power(signal):
@@ -1852,8 +1889,14 @@ def power(signal):
     if not isinstance(signal, pyfar.Signal):
         raise ValueError(f"signal is type '{signal.__class__}'"
                          " but must be of type 'Signal'.")
+
+    if isinstance(signal, (pyfar.Signal, pyfar.TimeData)):
+        if signal.complex:
+            raise ValueError((
+                "'power' is not implemented for complex time signals."))
+
     # return and compute data
-    return np.sum(signal.time**2, axis=-1)/signal.n_samples
+    return np.sum(np.abs(signal.time)**2, axis=-1)/signal.n_samples
 
 
 def rms(signal):
@@ -1889,6 +1932,11 @@ def rms(signal):
         raise ValueError(f"signal is type '{signal.__class__}'"
                          " but must be of type 'Signal'.")
 
+    if isinstance(signal, (pyfar.Signal, pyfar.TimeData)):
+        if signal.complex:
+            raise ValueError((
+                "'rms' is not implemented for complex time signals."))
+
     # return and compute data
     return np.sqrt(power(signal))
 
@@ -1973,15 +2021,23 @@ def average(signal, mode='linear', caxis=None, weights=None, keepdims=False,
                                pyfar.TimeData)):
         raise TypeError(("Input data has to be of type 'Signal', 'TimeData' "
                          "or 'FrequencyData'."))
+
+    if isinstance(signal, (pyfar.Signal, pyfar.TimeData)):
+        if signal.complex and mode == 'power':
+            raise ValueError((
+                "'power' is not implemented for complex time signals."))
+
     if type(signal) is pyfar.TimeData and mode in (
             'log_magnitude_zerophase', 'magnitude_zerophase',
             'magnitude_phase', 'power',):
         raise ValueError((
             f"mode is '{mode}' and signal is type '{signal.__class__}'"
             " but must be of type 'Signal' or 'FrequencyData'."))
+
     if nan_policy not in ('propagate', 'omit', 'raise'):
         raise ValueError("nan_policy has to be 'propagate', 'omit', or"
                          "'raise'.")
+
     # check for caxis
     if caxis and np.max(caxis) > len(signal.cshape):
         raise ValueError(('The maximum of caxis needs to be smaller than '
@@ -2043,9 +2099,11 @@ def average(signal, mode='linear', caxis=None, weights=None, keepdims=False,
     # return average data as pyfar object, depending on input signal type
     if isinstance(signal, pyfar.Signal):
         return pyfar.Signal(data, signal.sampling_rate, signal.n_samples,
-                            signal.domain, signal.fft_norm, signal.comment)
+                            signal.domain, signal.fft_norm, signal.comment,
+                            signal.complex)
     elif isinstance(signal, pyfar.TimeData):
-        return pyfar.TimeData(data, signal.times, signal.comment)
+        return pyfar.TimeData(data, signal.times, signal.comment,
+                              signal.complex)
     else:
         return pyfar.FrequencyData(data, signal.frequencies, signal.comment)
 
@@ -2216,6 +2274,10 @@ def normalize(signal, reference_method='max', domain='time',
         raise ValueError((
             f"domain is '{domain}' and signal is type '{signal.__class__}'"
             " but must be of type 'Signal' or 'FrequencyData'."))
+    if isinstance(signal, (pyfar.TimeData, pyfar.Signal)):
+        if signal.complex and reference_method in ['energy', 'power', 'rms']:
+            raise ValueError("'energy', 'power', and 'rms' reference method "
+                             "is not implemented for complex time signals.")
     if isinstance(limits, (int, float)) or len(limits) != 2:
         raise ValueError("limits must be an array like of length 2.")
     if tuple(limits) != (None, None) and \