Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mt_metadata \ mt_metadata \ features \ fc_coherence.py: 100%

1# ===================================================== 

2# Imports 

3# ===================================================== 

4from typing import Annotated 

5 

6import numpy as np 

7from pydantic import computed_field, Field, model_validator 

8 

9from mt_metadata.common.enumerations import StrEnumerationBase 

10from mt_metadata.features.coherence import Coherence 

11from mt_metadata.features.feature import Feature 

12 

13 

14# ===================================================== 

15class BandDefinitionTypeEnum(StrEnumerationBase): 

16 Q = "Q" 

17 fractional_bandwidth = "fractional bandwidth" 

18 user_defined = "user defined" 

19 

20 

21class QRadiusEnum(StrEnumerationBase): 

22 constant_Q = "constant Q" 

23 user_defined = "user defined" 

24 

25 

26class FCCoherence(Coherence, Feature): 

27 minimum_fcs: Annotated[ 

28 int, 

29 Field( 

30 default=2, 

31 description="The minimum number of Fourier coefficients needed to compute the feature.", 

32 alias=None, 

33 json_schema_extra={ 

34 "units": None, 

35 "required": True, 

36 "examples": ["2"], 

37 }, 

38 ), 

39 ] 

40 

41 band_definition_type: Annotated[ 

42 BandDefinitionTypeEnum, 

43 Field( 

44 default="Q", 

45 description="How the feature frequency bands are defined.", 

46 alias=None, 

47 json_schema_extra={ 

48 "units": None, 

49 "required": True, 

50 "examples": ["user defined"], 

51 }, 

52 ), 

53 ] 

54 

55 q_radius: Annotated[ 

56 QRadiusEnum, 

57 Field( 

58 default="constant Q", 

59 description="How the feature frequency bands are defined.", 

60 alias=None, 

61 json_schema_extra={ 

62 "units": None, 

63 "required": True, 

64 "examples": ["user defined"], 

65 }, 

66 ), 

67 ] 

68 

69 @model_validator(mode="before") 

70 @classmethod 

71 def set_defaults(cls, data: dict) -> dict: 

72 data["name"] = "fc_coherence" 

73 data["domain"] = "frequency" 

74 data["description"] = ( 

75 "Magnitude-squared coherence computed from frequency-domain Fourier coefficients (FCs). " 

76 "Cxy(f) = |Sxy(f)|^2 / (Sxx(f) * Syy(f)), where Sxy is the cross-power spectrum, " 

77 "Sxx and Syy are auto-power spectra, all estimated by averaging over windows." 

78 ) 

79 return data 

80 

81 @computed_field 

82 @property 

83 def channel_pair_str(self) -> str: 

84 return f"{self.channel_1}, {self.channel_2}" 

85 

86 def compute( 

87 self, fc1: np.ndarray, fc2: np.ndarray 

88 ) -> tuple[np.ndarray | None, np.ndarray]: 

89 """ 

90 Compute magnitude-squared coherence from FCs. 

91 

92 Parameters 

93 ---------- 

94 fc1 : np.ndarray 

95 Fourier coefficients for channel 1, shape (n_windows, n_freqs) 

96 fc2 : np.ndarray 

97 Fourier coefficients for channel 2, shape (n_windows, n_freqs) 

98 

99 Returns 

100 ------- 

101 freqs : np.ndarray 

102 Frequency axis (if available, else None) 

103 coherence : np.ndarray 

104 Magnitude-squared coherence, shape (n_freqs,) 

105 """ 

106 # Cross-power and auto-powers 

107 sxy = np.mean(fc1 * np.conj(fc2), axis=0) 

108 sxx = np.mean(np.abs(fc1) ** 2, axis=0) 

109 syy = np.mean(np.abs(fc2) ** 2, axis=0) 

110 

111 # Magnitude-squared coherence with protection against division by zero 

112 denominator = sxx * syy 

113 

114 # Use numpy error handling to suppress division warnings 

115 with np.errstate(divide="ignore", invalid="ignore"): 

116 coherence = np.abs(sxy) ** 2 / denominator 

117 

118 # Replace any infinite or NaN values with 0 

119 coherence = np.where(np.isfinite(coherence), coherence, 0.0) 

120 return None, coherence