Add frequency-domain coherence estimation with resolution mapping

cmvdr/data_gen/f0_manager.py

@@ -648,20 +648,72 @@ def compute_harmonic_and_modulation_sets_global_coherence(cls, sig, harmonic_fre
                                   max_freq_cyclic_hz=cfg_cyc['freq_range_cyclic'][1])
 
         max_bin = -1
-        if harmonic_freqs_est.size > 0:
-            max_bin = int(np.ceil((3 * SFT.delta_f + np.max(np.abs(harmonic_freqs_est))) / SFT.delta_f))
-        rho = CoherenceManager.compute_coherence(sig, SFT, mod_coherence, max_bin, min_relative_power=1.e+3)
+        # if harmonic_freqs_est.size > 0:
+        #     max_bin = int(np.ceil((3 * SFT.delta_f + np.max(np.abs(harmonic_freqs_est))) / SFT.delta_f))
+
+        # Choose coherence computation method based on config
+        use_freq_domain = cfg_cyc.get('use_freq_domain_coherence', False)
+
+        if use_freq_domain:
+            # Frequency-domain coherence path
+            use_stft = cfg_cyc.get('freq_coherence_stft_enabled', False)
+            interpolation = cfg_cyc.get('freq_coherence_interpolation', 'none')
+            apply_phase_correction = cfg_cyc.get('freq_coherence_apply_phase_correction', True)
+
+            # Create high-res STFT if specified
+            nfft_coherence = cfg_cyc.get('freq_coherence_nfft', None)
+            if nfft_coherence is not None and use_stft:
+                # Create a high-resolution STFT object
+                from scipy.signal import ShortTimeFFT, get_window
+                # Use 'hann' window like the standard STFT (or extract from config if needed)
+                win_coherence = get_window('hann', nfft_coherence)
+                # Use same hop as original, or scale proportionally
+                hop_coherence = SFT.hop
+                SFT_coherence = ShortTimeFFT(win_coherence, hop_coherence, fs=SFT.fs, 
+                                            mfft=nfft_coherence, scale_to=SFT.scaling)
+            else:
+                SFT_coherence = SFT
+
+            rho = CoherenceManager.compute_coherence_freq_shifted(
+                sig, SFT_coherence, mod_coherence.alpha_vec_hz_, max_bin,
+                min_relative_power=1.e+3,
+                use_stft=use_stft,
+                interpolation=interpolation,
+                apply_phase_correction=apply_phase_correction
+            )
+        else:
+            # Time-domain modulation path (original)
+            rho = CoherenceManager.compute_coherence(sig, SFT, mod_coherence, max_bin, min_relative_power=1.e+3)
+
         if 0:
             cc0 = np.where(mod_coherence.alpha_vec_hz_ == 0)[0][0]
             rho_no0 = np.delete(rho, cc0, axis=0)
             alpha_no0 = np.delete(mod_coherence.alpha_vec_hz_, cc0, axis=0)
             CoherenceManager.plot_coherence_matrix(rho_no0, alpha_no0, SFT)
 
         # retain highly coherent modulated components only
-        harm_info = CoherenceManager.calculate_harmonic_info_from_coherence(mod_coherence.alpha_vec_hz_, rho,
-                                                                            thr=cfg_cyc['harmonic_threshold'],
-                                                                            P_max_cfg=cfg_cyc['P_max'],
-                                                                            nfft_real=SFT.mfft // 2 + 1)
+        # Determine if frequency resolution mapping is needed
+        if use_freq and (nfft_coherence is not None and use_stft):
+            # Using high-res STFT for coherence
+            delta_f_coherence = SFT_coherence.delta_f
+            delta_f_beamforming = SFT.delta_f
+        elif use_freq and not use_stft:
+            # Using full-file DFT for coherence
+            delta_f_coherence = SFT.fs / len(sig)
+            delta_f_beamforming = SFT.delta_f
+        else:
+            # Time-domain method - coherence computed at beamforming resolution
+            delta_f_coherence = None
+            delta_f_beamforming = None
+
+        harm_info = CoherenceManager.calculate_harmonic_info_from_coherence(
+            mod_coherence.alpha_vec_hz_, rho,
+            thr=cfg_cyc['harmonic_threshold'],
+            P_max_cfg=cfg_cyc['P_max'],
+            nfft_real=SFT.mfft // 2 + 1,
+            delta_f_coherence=delta_f_coherence,
+            delta_f_beamforming=delta_f_beamforming
+        )
         mod_amount = F0ChangeAmount.small
 
         return harm_info, mod_amount