Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mth5 \ mth5 \ clients \ nims.py: 93%

1#!/usr/bin/env python 

2# coding: utf-8 

3 

4# ============================================================================= 

5# Imports 

6# ============================================================================= 

7from pathlib import Path 

8 

9from mth5 import read_file 

10from mth5.clients.base import ClientBase 

11from mth5.io.nims import NIMSCollection 

12from mth5.mth5 import MTH5 

13 

14 

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

16 

17 

18class NIMSClient(ClientBase): 

19 def __init__( 

20 self, 

21 data_path: str | Path, 

22 sample_rates: list[int] = [1, 8], 

23 save_path: str | Path | None = None, 

24 calibration_path: str | Path | None = None, 

25 mth5_filename: str = "from_nims.h5", 

26 **kwargs, 

27 ): 

28 super().__init__( 

29 data_path, 

30 save_path=save_path, 

31 sample_rates=sample_rates, 

32 mth5_filename=mth5_filename, 

33 **kwargs, 

34 ) 

35 

36 self.calibration_path = calibration_path 

37 self.collection = NIMSCollection(self.data_path) 

38 

39 @property 

40 def calibration_path(self) -> Path | None: 

41 """ 

42 Path to calibration data. 

43 

44 Returns 

45 ------- 

46 Path or None 

47 Path to calibration file, or None if not set. 

48 

49 Examples 

50 -------- 

51 >>> client = NIMSClient('data_dir') 

52 >>> client.calibration_path = 'calib.dat' 

53 >>> print(client.calibration_path) 

54 PosixPath('calib.dat') 

55 """ 

56 return self._calibration_path 

57 

58 @calibration_path.setter 

59 def calibration_path(self, value: str | Path | None) -> None: 

60 """ 

61 Set the calibration path. 

62 

63 Parameters 

64 ---------- 

65 value : str or Path or None 

66 Path to calibration file, or None. 

67 

68 Raises 

69 ------ 

70 IOError 

71 If the path does not exist. 

72 

73 Examples 

74 -------- 

75 >>> client = NIMSClient('data_dir') 

76 >>> client.calibration_path = 'calib.dat' 

77 """ 

78 if value is not None: 

79 self._calibration_path = Path(value) 

80 if not self._calibration_path.exists(): 

81 raise IOError(f"Could not find {self._calibration_path}") 

82 else: 

83 self._calibration_path = None 

84 

85 def get_run_dict(self) -> dict: 

86 """ 

87 Get run information from the NIMS collection. 

88 

89 Returns 

90 ------- 

91 dict 

92 Dictionary of run information. 

93 

94 Examples 

95 -------- 

96 >>> client = NIMSClient('data_dir') 

97 >>> runs = client.get_run_dict() 

98 >>> print(list(runs.keys())) 

99 ['station1', 'station2'] 

100 """ 

101 return self.collection.get_runs( 

102 sample_rates=self.sample_rates, 

103 calibration_path=self.calibration_path, 

104 ) 

105 

106 def get_survey(self, station_dict: dict) -> str: 

107 """ 

108 Get survey name from a dictionary of a single station of runs. 

109 

110 Parameters 

111 ---------- 

112 station_dict : dict 

113 Dictionary of runs for a station. 

114 

115 Returns 

116 ------- 

117 str 

118 Survey name. 

119 

120 Examples 

121 -------- 

122 >>> client = NIMSClient('data_dir') 

123 >>> runs = client.get_run_dict() 

124 >>> survey = client.get_survey(runs['station1']) 

125 >>> print(survey) 

126 'survey_name' 

127 """ 

128 return list( 

129 set([station_dict[k].survey.unique()[0] for k in station_dict.keys()]) 

130 )[0] 

131 

132 def make_mth5_from_nims( 

133 self, 

134 survey_id: str = "default_survey", 

135 combine: bool = True, 

136 **kwargs, 

137 ) -> str | Path: 

138 """ 

139 Make an MTH5 file from Phoenix NIMS files. Splits into runs, accounts for filters. 

140 

141 Parameters 

142 ---------- 

143 survey_id : str, optional 

144 Survey identifier. Default is "default_survey". 

145 combine : bool, optional 

146 Whether to combine runs. Default is True. 

147 **kwargs 

148 Additional keyword arguments to set as attributes. 

149 

150 Returns 

151 ------- 

152 str or Path 

153 Path to the saved MTH5 file. 

154 

155 Examples 

156 -------- 

157 >>> client = NIMSClient('data_dir') 

158 >>> mth5_path = client.make_mth5_from_nims(survey_id='survey1') 

159 >>> print(mth5_path) 

160 'output_dir/from_nims.h5' 

161 """ 

162 for key, value in kwargs.items(): 

163 if value is not None: 

164 setattr(self, key, value) 

165 

166 if survey_id is None: 

167 self.collection.survey_id = "none" 

168 else: 

169 self.collection.survey_id = survey_id 

170 

171 runs = self.get_run_dict() 

172 

173 with MTH5(**self.h5_kwargs) as m: 

174 m.open_mth5(self.save_path, self.mth5_file_mode) 

175 survey_group = m.add_survey(self.collection.survey_id) 

176 

177 for station_id in runs.keys(): 

178 station_group = survey_group.stations_group.add_station(station_id) 

179 for run_id, run_df in runs[station_id].items(): 

180 run_group = station_group.add_run(run_id) 

181 # use iloc to avoid KeyError if the DataFrame index does not start at 0 

182 run_ts = read_file(run_df["fn"].iloc[0]) 

183 run_ts.run_metadata.id = run_id 

184 run_group.from_runts(run_ts) 

185 station_group.metadata.update(run_ts.station_metadata) 

186 station_group.write_metadata() 

187 

188 # update survey metadata from input station 

189 survey_group.update_metadata() 

190 

191 return self.save_path