Skip to content

Commit

Permalink
add adloc
Browse files Browse the repository at this point in the history
  • Loading branch information
zhuwq0 committed Nov 24, 2023
1 parent e9cd1e4 commit bd254e4
Show file tree
Hide file tree
Showing 3 changed files with 135 additions and 94 deletions.
74 changes: 74 additions & 0 deletions slurm/plot_catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,15 @@
gamma_catalog = pd.read_csv(gamma_file, parse_dates=["time"])
# gamma_catalog["depth_km"] = gamma_catalog["depth(m)"] / 1e3


# %%
adloc_file = f"{root_path}/{region}/adloc/adloc_events.csv"
adloc_exist = False
if os.path.exists(adloc_file):
adloc_exist = True
adloc_catalog = pd.read_csv(adloc_file, parse_dates=["time"])
# gamma_catalog["depth_km"] = gamma_catalog["depth(m)"] / 1e3

# %%
hypodd_file = f"{root_path}/{region}/hypodd/hypodd_ct.reloc"
hypodd_ct_exist = False
Expand Down Expand Up @@ -311,11 +320,24 @@
s=min(2, size_factor / len(gamma_catalog)),
alpha=1.0,
linewidth=0,
label=f"GaMMA: {len(gamma_catalog)}",
)
ax[0, 1].set_title(f"GaMMA: {len(gamma_catalog)}")
# xlim = ax[0, 1].get_xlim()
# ylim = ax[0, 1].get_ylim()

if adloc_exist and (len(adloc_catalog) > 0):
ax[0, 1].scatter(
adloc_catalog["longitude"],
adloc_catalog["latitude"],
s=min(2, size_factor / len(adloc_catalog)),
alpha=1.0,
linewidth=0,
label=f"AdLoc: {len(adloc_catalog)}",
)
ax[0, 1].legend()
# ax[0, 1].set_title(f"AdLoc: {len(adloc_catalog)}")

if hypodd_ct_exist and (len(catalog_ct_hypodd) > 0):
ax[1, 0].scatter(
catalog_ct_hypodd["LON"],
Expand Down Expand Up @@ -378,6 +400,7 @@
vmin=cmin,
vmax=cmax,
cmap="viridis_r",
label=f"GaMMA: {len(gamma_catalog)}",
)
ax[0, 0].set_title(f"GaMMA: {len(gamma_catalog)}")
ax[0, 0].invert_yaxis()
Expand All @@ -386,6 +409,24 @@
else:
xlim = None
ylim = None

if adloc_exist and (len(adloc_catalog) > 0):
ax[0, 0].scatter(
adloc_catalog["longitude"],
adloc_catalog["depth_km"],
# c=adloc_catalog["depth_km"],
s=8000 / len(adloc_catalog),
alpha=1.0,
linewidth=0,
# vmin=cmin,
# vmax=cmax,
# cmap="viridis_r",
label=f"AdLoc: {len(adloc_catalog)}",
)
ax[0, 0].legend()
# ax[0, 0].set_title(f"AdLoc: {len(adloc_catalog)}")
ax[1, 0].set_xlim(xlim)
ax[1, 0].set_ylim(ylim)
if hypodd_ct_exist and (len(catalog_ct_hypodd) > 0):
ax[1, 0].scatter(
catalog_ct_hypodd["LON"],
Expand Down Expand Up @@ -458,6 +499,7 @@
vmin=cmin,
vmax=cmax,
cmap="viridis_r",
label=f"GaMMA: {len(gamma_catalog)}",
)
ax[0, 1].set_title(f"GaMMA: {len(gamma_catalog)}")
ax[0, 1].invert_yaxis()
Expand All @@ -466,6 +508,23 @@
else:
xlim = None
ylim = None
if adloc_exist and (len(adloc_catalog) > 0):
ax[0, 1].scatter(
adloc_catalog["latitude"],
adloc_catalog["depth_km"],
# c=adloc_catalog["depth_km"],
s=8000 / len(adloc_catalog),
alpha=1.0,
linewidth=0,
# vmin=cmin,
# vmax=cmax,
# cmap="viridis_r",
label=f"AdLoc: {len(adloc_catalog)}",
)
ax[0, 1].legend()
# ax[0, 1].set_title(f"AdLoc: {len(adloc_catalog)}")
ax[0, 1].set_xlim(xlim)
ax[0, 1].set_ylim(ylim)
if hypodd_ct_exist and (len(catalog_ct_hypodd) > 0):
ax[1, 1].scatter(
catalog_ct_hypodd["LAT"],
Expand Down Expand Up @@ -534,6 +593,9 @@
if gamma_exist:
ax[0, 0].plot(gamma_catalog["time"], gamma_catalog["magnitude"], "o", markersize=2, alpha=0.5, label="GaMMA")
ax[1, 0].plot(gamma_catalog["time"], gamma_catalog["magnitude"], "o", markersize=2, alpha=0.5, label="GaMMA")
if adloc_exist:
ax[0, 0].plot(adloc_catalog["time"], adloc_catalog["magnitude"], "o", markersize=2, alpha=0.5, label="AdLoc")
ax[1, 0].plot(adloc_catalog["time"], adloc_catalog["magnitude"], "o", markersize=2, alpha=0.5, label="AdLoc")
if hypodd_ct_exist:
ax[0, 0].plot(
catalog_ct_hypodd["time"], catalog_ct_hypodd["MAG"], "o", markersize=2, alpha=0.5, label="HypoDD (CT)"
Expand Down Expand Up @@ -562,6 +624,9 @@
if gamma_exist:
ax[0, 0].hist(gamma_catalog["magnitude"], bins=bins, alpha=0.5, label="GaMMA")
ax[1, 0].hist(gamma_catalog["magnitude"], bins=bins, alpha=0.5, label="GaMMA")
if adloc_exist:
ax[0, 0].hist(adloc_catalog["magnitude"], bins=bins, alpha=0.5, label="AdLoc")
ax[1, 0].hist(adloc_catalog["magnitude"], bins=bins, alpha=0.5, label="AdLoc")
if hypodd_ct_exist:
ax[0, 0].hist(catalog_ct_hypodd["MAG"], bins=bins, alpha=0.5, label="HypoDD (CT)")
if hypodd_cc_exist:
Expand Down Expand Up @@ -641,6 +706,15 @@ def plot3d(x, y, z, config, fig_name):
f"{root_path}/{result_path}/earthquake_location_gamma.html",
)

if adloc_exist and len(adloc_catalog) > 0:
plot3d(
adloc_catalog["longitude"],
adloc_catalog["latitude"],
# gamma_catalog["depth(m)"] / 1e3,
adloc_catalog["depth_km"],
config_plot3d,
f"{root_path}/{result_path}/earthquake_location_adloc.html",
)
if hypodd_ct_exist and len(catalog_ct_hypodd) > 0:
plot3d(
catalog_ct_hypodd["LON"],
Expand Down
95 changes: 29 additions & 66 deletions slurm/run_adloc.py
Original file line number Diff line number Diff line change
@@ -1,76 +1,39 @@
from typing import Dict, List, NamedTuple
import argparse
import os
import sys
from glob import glob
from pathlib import Path

from kfp import dsl
import torch


@dsl.component(base_image="zhuwq0/quakeflow:latest")
def run_gamma(
root_path: str,
region: str,
config: Dict,
index: int = 0,
picks_csv: str = "picks.csv",
events_csv: str = "events.csv",
protocol: str = "file",
bucket: str = "",
token: Dict = None,
) -> NamedTuple("outputs", events=str, picks=str):
import json
import os
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--root_path", type=str, default="local", help="root path")
parser.add_argument("--region", type=str, default="demo", help="region")
return parser.parse_args()

import fsspec
import numpy as np
import pandas as pd
from pyproj import Proj

# %%
fs = fsspec.filesystem(protocol=protocol, token=token)
args = parse_args()

# %%
result_path = f"{region}/gamma"
if not os.path.exists(f"{root_path}/{result_path}"):
os.makedirs(f"{root_path}/{result_path}")
# %%
root_path = args.root_path
region = args.region

# %%
station_json = f"{region}/obspy/stations.json"
# gamma_events_csv = f"{root_path}/gamma/gamma_events_{index:03d}.csv"
# gamma_picks_csv = f"{root_path}/gamma/gamma_picks_{index:03d}.csv"
result_path = f"{region}/adloc"
if not os.path.exists(f"{root_path}/{result_path}"):
os.makedirs(f"{root_path}/{result_path}")

## read picks
if protocol == "file":
picks = pd.read_csv(f"{root_path}/{picks_csv}")
else:
picks = pd.read_csv(f"{protocol}://{bucket}/{picks_csv}")
batch = 100

## read stations
if protocol == "file":
stations = pd.read_json(f"{root_path}/{station_json}", orient="index")
else:
with fs.open(f"{bucket}/{station_json}", "r") as fp:
stations = pd.read_json(fp, orient="index")
base_cmd = f"../ADLoc/run.py --config {root_path}/{region}/config.json --stations {root_path}/{region}/obspy/stations.json --events {root_path}/{region}/gamma/gamma_events.csv --picks {root_path}/{region}/gamma/gamma_picks.csv --result_path {root_path}/{region}/adloc --batch_size {batch}"
os.system(f"python {base_cmd} --device=cpu --epochs=1")


if __name__ == "__main__":
# import fire

# fire.Fire(run_gamma)

import json
import os
import sys

root_path = "local"
region = "demo"
if len(sys.argv) > 1:
root_path = sys.argv[1]
region = sys.argv[2]
with open(f"{root_path}/{region}/config.json", "r") as fp:
config = json.load(fp)

run_gamma.python_func(
root_path,
region=region,
config=config,
picks_csv=f"{region}/gamma/gamma_picks.csv",
events_csv=f"{region}/gamma/gamma_events.csv",
)
# num_gpu = torch.cuda.device_count()
# if num_gpu == 0:
# if os.uname().sysname == "Darwin":
# os.system(f"python {base_cmd} --device=mps")
# else:
# os.system(f"python {base_cmd} --device=cpu")
# else:
# os.system(f"torchrun --standalone --nproc_per_node {num_gpu} {base_cmd}")
60 changes: 32 additions & 28 deletions slurm/run_cctorch.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# %%
import argparse
import os
import sys
from glob import glob
from pathlib import Path

Expand All @@ -9,53 +10,56 @@

def parse_args():
parser = argparse.ArgumentParser()
## true
parser.add_argument("--root_path", type=str, default="local", help="root path")
parser.add_argument("--region", type=str, default="demo", help="region")
parser.add_argument("--dtct_pair", action="store_true", help="run convert_dtcc.py")
return parser.parse_args()


args = parse_args()

# %%
root_path = "local"
region = "demo"
root_path = args.root_path
region = args.region

output_path = f"{region}/templates"
if not os.path.exists(f"{root_path}/{output_path}"):
os.makedirs(f"{root_path}/{output_path}")
result_path = f"{region}/templates"
if not os.path.exists(f"{root_path}/{result_path}"):
os.makedirs(f"{root_path}/{result_path}")

mseed_list = glob(f"{root_path}/{region}/waveforms/????-???/??/")
mseed_list = [x + "*.mseed" for x in mseed_list]
with open(f"{root_path}/{region}/cctorch/mseed_list.txt", "w") as fp:
fp.write("\n".join(mseed_list))

batch = 2048
## based on GPU memory
# batch = 2048
# block_size1 = 1000
# block_size2 = 1000

batch = 100
block_size1 = 10
block_size2 = 10

if args.dtct_pair:
dt_ct = Path("relocation/hypodd/dt.ct")
dt_ct = f"{root_path}/{region}/hypodd/dt.ct"
pair_list = f"{root_path}/{region}/hypodd/event_pairs.txt"
lines = []
with open(dt_ct, "r") as fp:
for line in fp:
if line.startswith("#"):
ev1, ev2 = line.split()[1:3]
if ev1 > ev2:
ev1, ev2 = ev2, ev1
lines.append(f"{ev1},{ev2}\n")

print(f"Number of pairs: {len(lines)}")
with open(output_path / "event_pair.txt", "w") as fp:
print(f"Number of pairs from hypodd dt.ct: {len(lines)}")
with open(f"{root_path}/{region}/hypodd/event_pairs.txt", "w") as fp:
fp.writelines(lines)

# %%
os.system(
"python ../CCTorch/run.py --pair_list=templates/event_pair.txt --data_path=templates/template.dat --data_format=memmap --config=templates/config.json --batch_size=512 --result_path=templates/ccpairs"
)
base_cmd = f"../CCTorch/run.py --pair_list={root_path}/{region}/hypodd/event_pairs.txt --data_path1={root_path}/{region}/cctorch/template.dat --data_format1=memmap --config={root_path}/{region}/cctorch/config.json --batch_size={batch} --block_size1={block_size1} --block_size2={block_size2} --result_path={root_path}/{region}/cctorch/ccpairs"

else:
num_gpu = torch.cuda.device_count()
# base_cmd = f"../CCTorch/run.py --data_list1={root_path}/{region}/cctorch/event_index.txt --data_path1={root_path}/{region}/cctorch/template.dat --data_format1=memmap --config={root_path}/{region}/cctorch/config.json --batch_size={batch} --result_path={root_path}/{region}/cctorch/ccpairs"
base_cmd = f"../CCTorch/run.py --pair_list={root_path}/{region}/cctorch/event_pairs.txt --data_path1={root_path}/{region}/cctorch/template.dat --data_format1=memmap --config={root_path}/{region}/cctorch/config.json --batch_size={batch} --result_path={root_path}/{region}/cctorch/ccpairs"
if num_gpu == 0:
if os.uname().sysname == "Darwin":
os.system(f"python {base_cmd} --device=mps")
else:
os.system(f"python {base_cmd} --device=cpu")
base_cmd = f"../CCTorch/run.py --pair_list={root_path}/{region}/cctorch/event_pairs.txt --data_path1={root_path}/{region}/cctorch/template.dat --data_format1=memmap --config={root_path}/{region}/cctorch/config.json --batch_size={batch} --block_size1={block_size1} --block_size2={block_size2} --result_path={root_path}/{region}/cctorch/ccpairs"
num_gpu = torch.cuda.device_count()
if num_gpu == 0:
if os.uname().sysname == "Darwin":
os.system(f"python {base_cmd} --device=mps")
else:
os.system(f"torchrun --standalone --nproc_per_node {num_gpu} {base_cmd}")
os.system(f"python {base_cmd} --device=cpu")
else:
os.system(f"torchrun --standalone --nproc_per_node {num_gpu} {base_cmd}")

0 comments on commit bd254e4

Please sign in to comment.