update download waveform

AI4EPS · Oct 4, 2024 · 871cf52 · 871cf52
1 parent e18aed3
commit 871cf52
Show file tree

Hide file tree

Showing 7 changed files with 488 additions and 5 deletions.
diff --git a/scripts/args.py b/scripts/args.py
@@ -3,8 +3,21 @@
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("root_path", nargs="?", type=str, default="local", help="root path")
-    parser.add_argument("region", nargs="?", type=str, default="demo", help="region")
+    parser.add_argument("--region", type=str, default="demo", help="region")
+    parser.add_argument("--root_path", type=str, default="local", help="root path")
+
+    ## Cloud
+    parser.add_argument("--protocol", type=str, default="file", help="protocol (file, gs, s3)")
+    parser.add_argument("--bucket", type=str, default=None, help="bucket name")
+    parser.add_argument("--token", type=str, default=None, help="token")
+
+    # parser.add_argument("--bucket", type=str, default="quakeflow_catalog", help="bucket name")
+    # parser.add_argument("--protocol", type=str, default="gs", help="protocol (file, gs, s3)")
+    # parser.add_argument("--token", type=str, default="application_default_credentials.json", help="token")
+
+    ## Parallel
+    parser.add_argument("--num_nodes", type=int, default=1, help="number of nodes")
+    parser.add_argument("--node_rank", type=int, default=0, help="node rank")
 
     ## CCTorch
     parser.add_argument("--dtct_pair", action="store_true", help="run convert_dtcc.py")

diff --git a/scripts/download_catalog.py b/scripts/download_catalog.py
@@ -181,8 +181,12 @@ def visulization(config, events=None, stations=None, fig_name="catalog.png"):
     args = parse_args()
     root_path = args.root_path
     region = args.region
+    protocol = args.protocol
+    bucket = args.bucket
+    token = args.token
+
 
     with open(f"{root_path}/{region}/config.json", "r") as fp:
         config = json.load(fp)
 
-    download_catalog(root_path, region=region, config=config, protocol="file", bucket="", token=None)
+    download_catalog(root_path, region=region, config=config, protocol=protocol, bucket=bucket, token=token)
diff --git a/scripts/download_station.py b/scripts/download_station.py
@@ -329,10 +329,13 @@ def visulization(config, events=None, stations=None, fig_name="catalog.png"):
     args = parse_args()
     root_path = args.root_path
     region = args.region
+    protocol = args.protocol
+    bucket = args.bucket
+    token = args.token
 
     with open(f"{root_path}/{region}/config.json", "r") as fp:
         config = json.load(fp)
 
     config.update(config["obspy"])
 
-    download_station(root_path=root_path, region=region, config=config, protocol="file", bucket="", token=None)
+    download_station(root_path=root_path, region=region, config=config, protocol=protocol, bucket=bucket, token=token)
diff --git a/scripts/download_waveform_v2.py b/scripts/download_waveform_v2.py
@@ -318,6 +318,9 @@ def download_waveform(
     args = parse_args()
     root_path = args.root_path
     region = args.region
+    protocol = args.protocol
+    bucket = args.bucket
+    token = args.token
 
     with open(f"{root_path}/{region}/config.json", "r") as fp:
         config = json.load(fp)