add das event dataset

Author: zhuwq0

phasenet/data_reader.py

@@ -269,7 +269,13 @@ def __init__(
             self.data_dir = kwargs["data_dir"]
         elif format == "das":
             with open(kwargs["data_list"], "r") as fp:
-                self.data_list = fp.readlines()
+                self.data_list = [line.rstrip("\n") for line in fp.readlines()]
+
+            self.num_data = len(self.data_list)
+            self.data_dir = kwargs["data_dir"]
+        elif format == "das_event":
+            with open(kwargs["data_list"], "r") as fp:
+                self.data_list = [line.rstrip("\n") for line in fp.readlines()]
             self.num_data = len(self.data_list)
             self.data_dir = kwargs["data_dir"]
         else:
@@ -553,6 +559,24 @@ def read_das(self, fname, sampling_rate=100, highpass_filter=0.0):
         }
         return meta
 
+    def read_das_event(self, fname):
+        with fsspec.open(fname, "rb", anon=True) as fp:
+            with h5py.File(fp, "r") as f:
+                data = f["data"][:]
+                begin_time = f["data"].attrs["begin_time"]
+
+        data = np.array(data).T  # nt, nx
+        nt, nx = data.shape
+        # tmp = np.zeros([nt, nx, 3], dtype=np.float32)
+        # tmp[:, :, -1] = data
+        tmp = np.repeat(data[:, :, np.newaxis], 3, axis=-1)
+        meta = {
+            "data": tmp,
+            "t0": begin_time,
+            "station_id": [f"{x}" for x in range(nx)],
+        }
+        return meta
+
     def read_sac(self, fname):
         mseed = obspy.read(fname)
         mseed = mseed.detrend("spline", order=2, dspline=5 * mseed[0].stats.sampling_rate)
@@ -960,6 +984,8 @@ def __getitem__(self, i):
             base_name = ""
         elif self.format == "das":
             meta = self.read_das(base_name, sampling_rate=self.sampling_rate, highpass_filter=self.highpass_filter)
+        elif self.format == "das_event":
+            meta = self.read_das_event(base_name)
         else:
             raise (f"{self.format} does not support!")
 

phasenet/predict.py

@@ -162,7 +162,7 @@ def pred_fn(args, data_reader, figure_dir=None, prob_dir=None, log_dir=None):
                 tmp = fname_batch[0].decode().split(",")[0].split("/")
                 # subdir = "/".join(tmp[-1 - 3 : -1])
                 subdir = "/".join(tmp[-1 - args.subdir_level : -1])
-                fname = tmp[-1].rstrip("\n").rstrip(".mseed").rstrip(".ms") + ".csv"
+                fname = tmp[-1].rstrip("\n").rstrip(".mseed").rstrip(".ms").rstrip(".h5") + ".csv"
                 # csv_name = f"quakeflow_catalog/NC/phasenet/{subdir}/{fname}"
                 # csv_name = f"quakeflow_catalog/SC/phasenet/{subdir}/{fname}"
                 if not os.path.exists(os.path.join(args.result_dir, "picks", subdir)):