estimate PS ratio

AI4EPS · Nov 28, 2024 · 1e0232d · 1e0232d
1 parent e110273
commit 1e0232d
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/scripts/run_event_association.py b/scripts/run_event_association.py
@@ -172,7 +172,7 @@ def associate(
         # (ts - tp) = (ps_ratio - 1) tp = tt * 2 * (ps_ratio - 1) / (ps_ratio + 1)
 
         event = event.sort_values(by="num_picks", ascending=True)
-        ps_delta = event["travel_time_s"].values * 2 * (VPVS_RATIO - 1) / (VPVS_RATIO + 1)
+        ps_delta = event["travel_time"].values * 2 * (VPVS_RATIO - 1) / (VPVS_RATIO + 1)
         # t1 = event["timestamp_center"].values - ps_delta * 1.1 - 1.0
         # t2 = event["timestamp_center"].values + ps_delta * 1.1 + 1.0
         t1 = event["timestamp_center"].values - (ps_delta * 0.6 + 1.0)
@@ -183,9 +183,11 @@ def associate(
         # picks.loc[group_id, "event_index"] = np.where(
         #     mask.any(axis=0), index.values[mask.argmax(axis=0)], picks.loc[group_id, "event_index"]
         # )
-        picks.loc[group_id, "event_index"] = np.where(
-            mask.any(axis=0), event["event_index"].values[mask.argmax(axis=0)], -1
-        )
+        mask_true = mask.any(axis=0)
+        mask_idx = mask.argmax(axis=0)
+        picks.loc[group_id, "event_index"] = np.where(mask_true, event["event_index"].values[mask_idx], -1)
+        picks.loc[group_id, "sp_ratio"] = np.where(mask_true, event["sp_ratio"].values[mask_idx], np.nan)
+        picks.loc[group_id, "event_amplitude"] = np.where(mask_true, event["event_amplitude"].values[mask_idx], np.nan)
 
     picks.reset_index(inplace=True)
 

diff --git a/scripts/run_phasenet_plus.py b/scripts/run_phasenet_plus.py
@@ -73,8 +73,10 @@ def run_phasenet(
         processed = ["/".join(f.replace(".csv", "").split("/")[-subdir - 1 :]) for f in processed]
         processed = [p[:-1] for p in processed]  ## remove the channel suffix
         print(f"Number of processed files: {len(processed)}")
+        keys = sorted(list(set(mseed_3c.keys()) - set(processed)))
+    else:
+        keys = list(mseed_3c.keys())
 
-    keys = sorted(list(set(mseed_3c.keys()) - set(processed)))
     print(f"Number of unprocessed files: {len(keys)}")
     keys = list(np.array_split(keys, num_nodes)[node_rank])
     print(f"Node {node_rank:03d}/{num_nodes:03d}: processing {len(keys)} files")
@@ -117,7 +119,7 @@ def run_phasenet(
     with open(f"{root_path}/{region}/config.json", "r") as fp:
         config = json.load(fp)
 
-    run_phasenet(root_path=root_path, region=region, config=config)
+    run_phasenet(root_path=root_path, region=region, config=config, overwrite=args.overwrite)
 
     if num_nodes == 1:
         os.system(f"python merge_phasenet_plus_picks.py --region {region}")