update pick dicts into arguments style

Author: wanghaoyucn

eqnet/data/seismic_trace.py

@@ -398,6 +398,8 @@ def __init__(
         min_snr=3.0,
         stack_event=False,
         stack_strategy='same_sensor',
+        picks_dict=None,
+        events_dict=None,
         stack_noise=False,
         flip_polarity=False,
         drop_channel=False,
@@ -503,8 +505,10 @@ def __init__(
             elif self.data_path is not None:
                 self.base_dir = data_path
 
-            self.picks_dict = pd.read_csv(os.path.join(self.base_dir, "picks_train.csv"), usecols=['event_id', 'station_id', 'snr', 'phase_status', 'instrument'])
-            self.events_dict = pd.read_csv(os.path.join(self.base_dir, "events_train.csv"), usecols=['event_id', 'magnitude', 'event_time', 'depth_km'])
+            picks_dict = os.path.join(self.base_dir, "picks_train.csv") if (picks_dict is None) else picks_dict
+            events_dict = os.path.join(self.base_dir, "events_train.csv") if (events_dict is None) else events_dict
+            self.picks_dict = pd.read_csv(picks_dict, usecols=['event_id', 'station_id', 'snr', 'phase_status', 'instrument'])
+            self.events_dict = pd.read_csv(events_dict, usecols=['event_id', 'magnitude', 'event_time', 'depth_km'])
             self.events_dict['year'] = self.events_dict['event_time'].apply(lambda x: int(x[:4]))
             self.picks_dict['snr'] = self.picks_dict['snr'].apply(lambda x: np.array([float(number) for number in x.split()[1:-1]]))
             temp = self.picks_dict.groupby('event_id')['snr'].apply(lambda x: np.concatenate(x.values)).reset_index()

train.py

@@ -379,6 +379,8 @@ def main(args):
             format=args.format,  # ["h5", "hf"]
             training=True,
             stack_event=args.stack_event,
+            picks_dict=args.picks_dict,
+            events_dict=args.events_dict,
             stack_noise=args.stack_noise,
             flip_polarity=args.flip_polarity,
             drop_channel=args.drop_channel,
@@ -716,6 +718,8 @@ def get_args_parser(add_help=True):
     parser.add_argument("--test-label-list", default="+", type=None, help="test label path")
     parser.add_argument("--test-noise-list", default="+", type=None, help="test noise list")
     parser.add_argument("--test-hdf5-file", default=None, type=str, help="hdf5 file for testing")
+    parser.add_argument("--picks-dict", default=None, type=str, help="picks dictionary for training augmentation")
+    parser.add_argument("--events-dict", default=None, type=str, help="events dictionary for training augmentation")
     parser.add_argument("--dataset", default="", type=str, help="dataset name")
     parser.add_argument("--model", default="phasenet_das", type=str, help="model name")
     parser.add_argument("--backbone", default="unet", type=str, help="model backbone")