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

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

2# Imports 

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

4from typing import Annotated 

5 

6import numpy as np 

7import xarray as xr 

8from pydantic import Field, field_validator, PrivateAttr, ValidationInfo 

9 

10from mt_metadata.base import MetadataBase 

11from mt_metadata.common import Comment 

12from mt_metadata.common.enumerations import StrEnumerationBase 

13 

14 

15# ===================================================== 

16class DomainEnum(StrEnumerationBase): 

17 time = "time" 

18 frequency = "frequency" 

19 fc = "fc" 

20 ts = "ts" 

21 fourier = "fourier" 

22 

23 

24class Feature(MetadataBase): 

25 name: Annotated[ 

26 str, 

27 Field( 

28 default="", 

29 description="Name of the feature.", 

30 alias=None, 

31 json_schema_extra={ 

32 "units": None, 

33 "required": True, 

34 "examples": ["simple coherence"], 

35 }, 

36 ), 

37 ] 

38 

39 description: Annotated[ 

40 str, 

41 Field( 

42 default="", 

43 description="A full description of what the feature estimates.", 

44 alias=None, 

45 json_schema_extra={ 

46 "units": None, 

47 "required": True, 

48 "examples": [ 

49 "Simple coherence measures the coherence between measured electric and magnetic fields." 

50 ], 

51 }, 

52 ), 

53 ] 

54 

55 domain: Annotated[ 

56 DomainEnum, 

57 Field( 

58 default=DomainEnum.frequency, 

59 description="Temporal domain the feature is estimated in [ 'frequency' | 'time' ]", 

60 alias=None, 

61 json_schema_extra={ 

62 "units": None, 

63 "required": True, 

64 "examples": ["frequency"], 

65 }, 

66 ), 

67 ] 

68 

69 comments: Annotated[ 

70 Comment, 

71 Field( 

72 default_factory=lambda: Comment(), # type: ignore 

73 description="Any comments about the feature", 

74 alias=None, 

75 json_schema_extra={ 

76 "units": None, 

77 "required": False, 

78 "examples": ["estimated using hilburt transform."], 

79 }, 

80 ), 

81 ] 

82 

83 data: Annotated[ 

84 xr.DataArray | xr.Dataset | np.ndarray | None, 

85 Field( 

86 default=None, 

87 description="The data associated with the feature.", 

88 alias=None, 

89 json_schema_extra={ 

90 "units": None, 

91 "required": True, 

92 "examples": ["path/to/datafile.nc"], 

93 }, 

94 ), 

95 ] 

96 

97 _supported_features: dict[str, type] = PrivateAttr(default_factory=dict) 

98 

99 @field_validator("comments", mode="before") 

100 @classmethod 

101 def validate_comments(cls, value, info: ValidationInfo) -> Comment: 

102 if isinstance(value, str): 

103 return Comment(value=value) 

104 return value 

105 

106 @field_validator("data", mode="before") 

107 @classmethod 

108 def validate_data( 

109 cls, value, info: ValidationInfo 

110 ) -> xr.DataArray | xr.Dataset | np.ndarray | None: 

111 if value is None: 

112 return None 

113 elif not isinstance(value, (xr.DataArray, xr.Dataset, np.ndarray)): 

114 raise TypeError("Data must be a numpy array, xarray, or None.") 

115 return value 

116 

117 @classmethod 

118 def from_feature_id(cls, meta_dict): 

119 """ 

120 Factory: instantiate the correct feature class based on 'feature_id'. 

121 

122 not sure this is needed anymore. 

123 """ 

124 if "feature_id" not in meta_dict: 

125 raise KeyError("Feature metadata must include 'feature_id'.") 

126 feature_id = meta_dict["feature_id"] 

127 

128 # Import here to avoid circular dependencies 

129 from mt_metadata.features.registry import SUPPORTED_FEATURE_DICT 

130 

131 supported = SUPPORTED_FEATURE_DICT 

132 

133 if feature_id not in supported: 

134 raise KeyError( 

135 f"Unknown feature_id '{feature_id}'. Supported: {list(supported.keys())}" 

136 ) 

137 feature_cls = supported[feature_id] 

138 obj = feature_cls() 

139 obj.from_dict(meta_dict) 

140 return obj