Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mt_metadata \ mt_metadata \ processing \ aurora \ stations.py: 94%

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

2# Imports 

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

4from typing import Annotated 

5 

6import pandas as pd 

7from pydantic import Field, field_validator, ValidationInfo 

8 

9from mt_metadata.base import MetadataBase 

10from mt_metadata.processing.aurora.station import Station 

11 

12 

13# ===================================================== 

14class Stations(MetadataBase): 

15 remote: Annotated[ 

16 list[Station], 

17 Field( 

18 default_factory=list, 

19 description="list of remote sites", 

20 alias=None, 

21 json_schema_extra={ 

22 "units": None, 

23 "required": True, 

24 "examples": ["10"], 

25 }, 

26 ), 

27 ] 

28 

29 local: Annotated[ 

30 Station, 

31 Field( 

32 default_factory=Station, # type: ignore 

33 description="local site", 

34 alias=None, 

35 json_schema_extra={ 

36 "units": None, 

37 "required": True, 

38 "examples": ["10"], 

39 }, 

40 ), 

41 ] 

42 

43 @field_validator("remote", mode="before") 

44 def validate_remote( 

45 cls, value: list[Station], info: ValidationInfo 

46 ) -> list[Station]: 

47 """ 

48 Method for unpacking rr_station info into mt_metadata object. 

49 

50 Developmnent Notes: 

51 This function was raising an exception when trying to populate an aurora.Processing object 

52 from a json.loads() dict. 

53 TODO: add a description of input variable and use cases, ... it seems that we may not want 

54 to support multiple rr stations yet. 

55 

56 Parameters 

57 ---------- 

58 rr_station 

59 

60 Returns 

61 ------- 

62 list of Station objects 

63 

64 """ 

65 remote = [] 

66 if isinstance(value, list): 

67 for item in value: 

68 if isinstance(item, Station): 

69 remote.append(item) 

70 elif isinstance(item, dict): 

71 try: 

72 station = Station() # type: ignore 

73 station.from_dict(item) 

74 remote.append(station) 

75 except Exception as e: 

76 raise ValueError("could not unpack dict to a Station object") 

77 else: 

78 raise TypeError( 

79 f"list item must be Station object not {type(item)}" 

80 ) 

81 

82 elif isinstance(value, dict): 

83 station = Station() 

84 station.from_dict(value) 

85 station.remote = True 

86 remote.append(station) 

87 

88 elif isinstance(value, Station): 

89 value.remote = True 

90 remote.append(value) 

91 

92 elif isinstance( 

93 value, str 

94 ): # TODO: Add doc; what is this doing? This does not affect self._remote. 

95 if len(value) > 4: 

96 raise ValueError(f"not sure to do with {type(value)}") 

97 # TODO: Add doc explaining what happens when rr_station is str of length 3. 

98 

99 else: 

100 raise ValueError(f"not sure to do with {type(value)}") 

101 

102 return remote 

103 

104 def add_remote(self, rr: Station | dict): 

105 """ 

106 add a remote station 

107 

108 Parameters 

109 ---------- 

110 rr: Station | dict 

111 remote station to add 

112 """ 

113 

114 if not isinstance(rr, (Station, dict)): 

115 raise TypeError(f"List entry must be a Station object not {type(rr)}") 

116 if isinstance(rr, dict): 

117 obj = Station() # type: ignore 

118 obj.from_dict(rr) 

119 

120 else: 

121 obj = rr 

122 

123 obj.remote = True 

124 

125 self.remote.append(obj) 

126 

127 @property 

128 def remote_dict(self) -> dict[str, Station]: 

129 """ 

130 need to have a dictionary, but it can't be an attribute cause that 

131 gets confusing when reading in a json file 

132 

133 Returns 

134 ------- 

135 dict[str, Station] 

136 dictionary of remote stations 

137 

138 """ 

139 return dict([(rr.id, rr) for rr in self.remote]) 

140 

141 def from_dataset_dataframe(self, df: pd.DataFrame): 

142 """ 

143 from a dataset dataframe 

144 

145 Parameters 

146 ---------- 

147 df: pd.DataFrame 

148 dataset dataframe to read from 

149 

150 Returns 

151 ------- 

152 None 

153 """ 

154 

155 # Handle empty DataFrame case 

156 if df.empty: 

157 return 

158 

159 station = df[df.remote == False].station.unique()[0] 

160 rr_stations = df[df.remote == True].station.unique() 

161 

162 self.local.from_dataset_dataframe(df[df.station == station]) 

163 

164 for rr_station in rr_stations: 

165 rr = Station() # type: ignore 

166 rr.from_dataset_dataframe(df[df.station == rr_station]) 

167 self.add_remote(rr) 

168 

169 def to_dataset_dataframe(self) -> pd.DataFrame: 

170 """ 

171 output a dataframe 

172 

173 Returns 

174 ------- 

175 pd.DataFrame 

176 dataframe representation of the station 

177 """ 

178 

179 df = self.local.to_dataset_dataframe() 

180 for rr in self.remote: 

181 remote_df = rr.to_dataset_dataframe() 

182 df = pd.concat([df, remote_df]) # , axis=1, ignore_index=True) 

183 

184 df.reset_index(inplace=True, drop=True) 

185 

186 return df 

187 

188 def get_station(self, station_id: str) -> Station: 

189 """ 

190 get a station object from the id 

191 

192 Parameters 

193 ---------- 

194 station_id: str 

195 ID of the station to retrieve 

196 

197 Returns 

198 ------- 

199 Station 

200 Station object corresponding to the given ID 

201 

202 """ 

203 

204 if self.local.id == station_id: 

205 return self.local 

206 

207 elif station_id in self.remote_dict.keys(): 

208 return self.remote_dict[station_id] 

209 

210 raise KeyError(f"could not find {station_id}")