Coverage for C: \ Users \ peaco \ OneDrive \ Documents \ GitHub \ mth5 \ mth5 \ clients \ metronix.py: 100%

1# -*- coding: utf-8 -*- 

2""" 

3Created on Wed Nov 27 11:23:50 2024 

4 

5@author: jpeacock 

6""" 

7 

8# ============================================================================= 

9# Imports 

10# ============================================================================= 

11import pandas as pd 

12 

13from mth5 import read_file 

14from mth5.clients.base import ClientBase 

15from mth5.io.metronix import MetronixCollection 

16from mth5.mth5 import MTH5 

17 

18 

19# ============================================================================= 

20 

21 

22class MetronixClient(ClientBase): 

23 def __init__( 

24 self, 

25 data_path, 

26 sample_rates=[128], 

27 save_path=None, 

28 calibration_path=None, 

29 mth5_filename="from_metronix.h5", 

30 **kwargs, 

31 ): 

32 super().__init__( 

33 data_path, 

34 save_path=save_path, 

35 sample_rates=sample_rates, 

36 mth5_filename=mth5_filename, 

37 **kwargs, 

38 ) 

39 self.calibration_path = calibration_path 

40 self.collection = MetronixCollection(self.data_path) 

41 

42 def get_run_dict(self, run_name_zeros=0): 

43 """ 

44 get run information 

45 

46 :return: DESCRIPTION 

47 :rtype: TYPE 

48 

49 """ 

50 

51 return self.collection.get_runs( 

52 sample_rates=self.sample_rates, 

53 run_name_zeros=run_name_zeros, 

54 calibration_path=self.calibration_path, 

55 ) 

56 

57 def get_survey_id(self, station_dict): 

58 """ 

59 get survey name from a dictionary of a single station of runs 

60 :param station_dict: DESCRIPTION 

61 :type station_dict: TYPE 

62 :return: DESCRIPTION 

63 :rtype: TYPE 

64 

65 """ 

66 

67 return list( 

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

69 )[0] 

70 

71 def set_station_metadata(self, station_dict, station_group): 

72 """ 

73 set station group metadata from information in the station dict 

74 

75 :param station_dict: DESCRIPTION 

76 :type station_dict: TYPE 

77 :param station_group: DESCRIPTION 

78 :type station_group: TYPE 

79 :return: DESCRIPTION 

80 :rtype: TYPE 

81 

82 """ 

83 

84 runs = pd.concat(station_dict.values()) 

85 station_group.metadata.location.latitude = runs.latitude.mean() 

86 station_group.metadata.location.longitude = runs.longitude.mean() 

87 station_group.metadata.location.elevation = runs.elevation.mean() 

88 station_group.write_metadata() 

89 

90 def make_mth5_from_metronix(self, run_name_zeros=0, **kwargs): 

91 """ 

92 Create an MTH5 from new ATSS + JSON style Metronix data. 

93 

94 :param **kwargs: DESCRIPTION 

95 :type **kwargs: TYPE 

96 :return: DESCRIPTION 

97 :rtype: TYPE 

98 

99 """ 

100 

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

102 if value is not None: 

103 setattr(self, key, value) 

104 

105 runs = self.get_run_dict(run_name_zeros=run_name_zeros) 

106 

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

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

109 

110 for station_id, station_dict in runs.items(): 

111 survey_id = self.get_survey_id(station_dict) 

112 survey_group = m.add_survey(survey_id) 

113 station_group = survey_group.stations_group.add_station(station_id) 

114 self.set_station_metadata(station_dict, station_group) 

115 

116 for run_id, run_df in station_dict.items(): 

117 run_group = station_group.add_run(run_id) 

118 for row in run_df.itertuples(): 

119 ch_ts = read_file(row.fn) 

120 run_group.from_channel_ts(ch_ts) 

121 run_group.update_metadata() 

122 station_group.update_metadata() 

123 survey_group.update_metadata 

124 

125 self.logger.info(f"Wrote MTH5 file to: {self.save_path}") 

126 

127 return self.save_path