Skip to content

Commit

Permalink
improve relative location
Browse files Browse the repository at this point in the history
  • Loading branch information
@zhuwq0
zhuwq0 committed Sep 3, 2024
1 parent 8e6ed1f commit 70b76d1
Show file tree
Hide file tree
Showing 6 changed files with 286 additions and 35 deletions.
8 changes: 6 additions & 2 deletions adloc/data.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,7 @@ def __init__(
events,
stations,
batch_size=1000,
valid_index=None,
config=None,
rank=0,
world_size=1,
Expand All @@ -213,6 +214,9 @@ def __init__(
self.stations = stations
self.config = config
self.batch_size = batch_size
if valid_index is None:
valid_index = np.ones(len(self.pairs), dtype=bool)
self.valid_index = valid_index

self.idx_batch = np.array_split(np.arange(len(self.pairs)), (len(self.pairs) - 1) // self.batch_size + 1)[
rank::world_size
Expand All @@ -226,14 +230,14 @@ def __len__(self):

def __getitem__(self, i):

idx = self.idx_batch[i]
idx = self.idx_batch[i][self.valid_index[self.idx_batch[i]]]
idx1_eve = self.pairs["event_index1"][idx]
idx2_eve = self.pairs["event_index2"][idx]
idx_eve = np.stack([idx1_eve, idx2_eve], axis=1)
idx_sta = self.pairs["station_index"][idx]
phase_weight = self.pairs["phase_score"][idx]
phase_type = self.pairs["phase_type"][idx]
phase_time = self.pairs["dd_time"][idx]
phase_time = self.pairs["phase_dtime"][idx]

return {
"idx_eve": torch.tensor(idx_eve, dtype=torch.long),
Expand Down
10 changes: 5 additions & 5 deletions docs/convert_dtcc.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@
station_index = []
phase_type = []
phase_score = []
dd_time = []
phase_dtime = []

stations.set_index("station_id", inplace=True)
events.set_index("event_index", inplace=True)
Expand All @@ -52,7 +52,7 @@
station_index.append(stations.loc[stid, "idx_sta"])
phase_type.append(mapping_phase_type_int[phase])
phase_score.append(weight)
dd_time.append(dt)
phase_dtime.append(dt)


dtypes = np.dtype(
Expand All @@ -62,21 +62,21 @@
("station_index", np.int32),
("phase_type", np.int32),
("phase_score", np.float32),
("dd_time", np.float32),
("phase_dtime", np.float32),
]
)
pairs_array = np.memmap(
os.path.join(result_path, "pair_dt.dat"),
mode="w+",
shape=(len(dd_time),),
shape=(len(phase_dtime),),
dtype=dtypes,
)
pairs_array["event_index1"] = event_index1
pairs_array["event_index2"] = event_index2
pairs_array["station_index"] = station_index
pairs_array["phase_type"] = phase_type
pairs_array["phase_score"] = phase_score
pairs_array["dd_time"] = dd_time
pairs_array["phase_dtime"] = phase_dtime
with open(os.path.join(result_path, "pair_dtypes.pkl"), "wb") as f:
pickle.dump(dtypes, f)

Expand Down
34 changes: 17 additions & 17 deletions docs/generate_pairs.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,10 +15,10 @@

def parse_args():
parser = argparse.ArgumentParser(description="Generate pairs")
parser.add_argument("--stations", type=str, default="test_data/synthetic/stations.csv")
parser.add_argument("--events", type=str, default="test_data/synthetic/events.csv")
parser.add_argument("--picks", type=str, default="test_data/synthetic/picks.csv")
parser.add_argument("--result_path", type=str, default="results/synthetic")
parser.add_argument("--stations", type=str, default="test_data/ridgecrest/stations.csv")
parser.add_argument("--events", type=str, default="test_data/ridgecrest/gamma_events.csv")
parser.add_argument("--picks", type=str, default="test_data/ridgecrest/gamma_picks.csv")
parser.add_argument("--result_path", type=str, default="results/ridgecrest")
return parser.parse_args()


Expand All @@ -36,7 +36,7 @@ def convert_dd(
event_index2 = []
phase_type = []
phase_score = []
dd_time = []
phase_dtime = []
for idx1, idx2 in tqdm(pairs, desc=f"CPU {i}", position=i):
picks1 = picks_by_event.get_group(idx1)
picks2 = picks_by_event.get_group(idx2)
Expand All @@ -52,15 +52,15 @@ def convert_dd(
station_index.extend(common["idx_sta"].values)
phase_type.extend(common["phase_type"].values)
phase_score.extend(common["phase_score"].values)
dd_time.extend(np.round(common["travel_time_x"].values - common["travel_time_y"].values, 5))
phase_dtime.extend(np.round(common["travel_time_x"].values - common["travel_time_y"].values, 5))

return {
"event_index1": event_index1,
"event_index2": event_index2,
"station_index": station_index,
"phase_type": phase_type,
"phase_score": phase_score,
"dd_time": dd_time,
"phase_dtime": phase_dtime,
}


Expand All @@ -73,11 +73,11 @@ def convert_dd(
os.makedirs(result_path)

# %%
MAX_PAIR_DIST = 5 # km
MAX_NEIGHBORS = 20
MAX_PAIR_DIST = 10 # km
MAX_NEIGHBORS = 50
MIN_NEIGHBORS = 8
MIN_OBS = 8
MAX_OBS = 20
MAX_OBS = 100
mapping_phase_type_int = {"P": 0, "S": 1}

# %%
Expand Down Expand Up @@ -191,16 +191,16 @@ def convert_dd(
station_index = np.concatenate([r["station_index"] for r in results])
phase_type = np.concatenate([r["phase_type"] for r in results])
phase_score = np.concatenate([r["phase_score"] for r in results])
dd_time = np.concatenate([r["dd_time"] for r in results])
phase_dtime = np.concatenate([r["phase_dtime"] for r in results])

# Filter large P and S time differences
idx = ((phase_type == 0) & (np.abs(dd_time) < 1.0)) | ((phase_type == 1) & (np.abs(dd_time) < 1.5))
idx = ((phase_type == 0) & (np.abs(phase_dtime) < 1.0)) | ((phase_type == 1) & (np.abs(phase_dtime) < 1.5))
event_index1 = event_index1[idx]
event_index2 = event_index2[idx]
station_index = station_index[idx]
phase_type = phase_type[idx]
phase_score = phase_score[idx]
dd_time = dd_time[idx]
phase_dtime = phase_dtime[idx]
print(f"Saving to disk: {len(event_index1)} pairs")
# np.savez_compressed(
# os.path.join(catalog_path, "adloc_dt.npz"),
Expand All @@ -209,7 +209,7 @@ def convert_dd(
# station_index=station_index,
# phase_type=phase_type,
# phase_score=phase_score,
# dd_time=dd_time,
# phase_dtime=phase_dtime,
# )

dtypes = np.dtype(
Expand All @@ -219,21 +219,21 @@ def convert_dd(
("station_index", np.int32),
("phase_type", np.int32),
("phase_score", np.float32),
("dd_time", np.float32),
("phase_dtime", np.float32),
]
)
pairs_array = np.memmap(
os.path.join(result_path, "pair_dt.dat"),
mode="w+",
shape=(len(dd_time),),
shape=(len(phase_dtime),),
dtype=dtypes,
)
pairs_array["event_index1"] = event_index1
pairs_array["event_index2"] = event_index2
pairs_array["station_index"] = station_index
pairs_array["phase_type"] = phase_type
pairs_array["phase_score"] = phase_score
pairs_array["dd_time"] = dd_time
pairs_array["phase_dtime"] = phase_dtime
with open(os.path.join(result_path, "pair_dtypes.pkl"), "wb") as f:
pickle.dump(dtypes, f)

Expand Down
Loading

0 comments on commit 70b76d1

Please sign in to comment.