diff --git a/adloc/adloc.py b/adloc/adloc.py index b9dd5dc..c90b6d2 100644 --- a/adloc/adloc.py +++ b/adloc/adloc.py @@ -189,7 +189,7 @@ def forward( resisudal[phase_type == type] = phase_time_ - t_ loss += torch.sum(F.huber_loss(t_, phase_time_, reduction="none") * phase_weight_) - return {"phase_time": pred_time, "residual_s": resisudal, "loss": loss} + return {"phase_time": pred_time, "residual": resisudal, "loss": loss} # %% diff --git a/adloc/sacloc2d.py b/adloc/sacloc2d.py index 49b58d9..8b131a4 100644 --- a/adloc/sacloc2d.py +++ b/adloc/sacloc2d.py @@ -178,6 +178,21 @@ def predict(self, X, event_index=0): else: return np.array([tt]).T + def weight(self, X, event_index=0): + """ + X: data_frame with columns ["timestamp", "x_km", "y_km", "z_km", "type"] + """ + station_index = X[:, 0].astype(int) + phase_type = X[:, 1].astype(int) + + if self.eikonal is None: + v = np.array([self.vel[t] for t in phase_type]) + tt = np.linalg.norm(self.events[event_index, :3] - self.stations[station_index, :3], axis=-1) / v + else: + tt = traveltime(event_index, station_index, phase_type, self.events, self.stations, self.eikonal) + + return np.array([tt]).T + def score(self, X, y=None, event_index=0): """ X: idx_sta, type, score, amp diff --git a/adloc/utils.py b/adloc/utils.py index 40c9f69..dac7f41 100644 --- a/adloc/utils.py +++ b/adloc/utils.py @@ -146,6 +146,8 @@ def is_data_valid(X, y): else: mean_residual_amp = 0.0 x, y, z, t = estimator.events[0] + if config["use_amplitude"]: + mag = estimator.magnitudes[0] event = { "idx_eve": event_index, ## inside adloc, idx_eve is used which starts from 0 to N events @@ -158,6 +160,7 @@ def is_data_valid(X, y): "num_picks": np.sum(inlier_mask), } if config["use_amplitude"]: + event["magnitude"] = mag event["adloc_residual_amplitude"] = mean_residual_amp event = pd.DataFrame([event]) else: diff --git a/app.py b/app.py index 70a8c5b..6e01c83 100644 --- a/app.py +++ b/app.py @@ -162,7 +162,7 @@ def run_adloc(picks, stations, config_): events = events.drop(["idx_eve", "x_km", "y_km", "z_km"], axis=1, errors="ignore") events.sort_values(["time"], inplace=True) - picks.rename({"mask": "adloc_mask", "residual_s": "adloc_residual_s"}, axis=1, inplace=True) + picks.rename({"mask": "adloc_mask", "residual_time": "adloc_residual_time", "residual_amplitude": "adloc_residual_amplitude"}, axis=1, inplace=True) picks["phase_type"] = picks["phase_type"].map({0: "P", 1: "S"}) picks = picks.drop(["idx_eve", "idx_sta"], axis=1, errors="ignore") picks.sort_values(["phase_time"], inplace=True) diff --git a/docs/run_adloc.py b/docs/run_adloc.py index ddda49a..7b626e1 100644 --- a/docs/run_adloc.py +++ b/docs/run_adloc.py @@ -298,16 +298,16 @@ for i in range(MAX_SST_ITER): with torch.inference_mode(): - picks["residual_s"] = ( - travel_time(idx_sta, idx_eve, phase_type, phase_time, phase_weight)["residual_s"].detach().numpy() + picks["residual"] = ( + travel_time(idx_sta, idx_eve, phase_type, phase_time, phase_weight)["residual"].detach().numpy() ) - # station_term = picks.groupby("idx_sta").agg({"residual_s": "mean"}).reset_index() + # station_term = picks.groupby("idx_sta").agg({"residual": "mean"}).reset_index() station_term = ( picks.groupby("idx_sta") - .apply(lambda x: weighted_mean(x["residual_s"], x["phase_score"])) - .reset_index(name="residual_s") + .apply(lambda x: weighted_mean(x["residual"], x["phase_score"])) + .reset_index(name="residual") ) - stations["station_term"] += stations["idx_sta"].map(station_term.set_index("idx_sta")["residual_s"]).fillna(0) + stations["station_term"] += stations["idx_sta"].map(station_term.set_index("idx_sta")["residual"]).fillna(0) travel_time.station_dt.weight.data = torch.tensor(stations["station_term"].values, dtype=torch.float32).view( -1, 1 ) diff --git a/docs/run_adloc_backup.py b/docs/run_adloc_backup.py index 585530f..5e55011 100644 --- a/docs/run_adloc_backup.py +++ b/docs/run_adloc_backup.py @@ -287,16 +287,16 @@ def main(args): weighted_mean = lambda x, w: np.sum(x * w) / np.sum(w) for i in range(MAX_SST_ITER): with torch.inference_mode(): - picks["residual_s"] = ( - travel_time(idx_sta, idx_eve, phase_type, phase_time, phase_weight)["residual_s"].detach().numpy() + picks["residual"] = ( + travel_time(idx_sta, idx_eve, phase_type, phase_time, phase_weight)["residual"].detach().numpy() ) - # station_term = picks.groupby("idx_sta").agg({"residual_s": "mean"}).reset_index() + # station_term = picks.groupby("idx_sta").agg({"residual": "mean"}).reset_index() station_term = ( picks.groupby("idx_sta") - .apply(lambda x: weighted_mean(x["residual_s"], x["phase_score"])) - .reset_index(name="residual_s") + .apply(lambda x: weighted_mean(x["residual"], x["phase_score"])) + .reset_index(name="residual") ) - stations["station_term"] += stations["idx_sta"].map(station_term.set_index("idx_sta")["residual_s"]).fillna(0) + stations["station_term"] += stations["idx_sta"].map(station_term.set_index("idx_sta")["residual"]).fillna(0) travel_time.station_dt.weight.data = torch.tensor(stations["station_term"].values, dtype=torch.float32).view( -1, 1 ) diff --git a/docs/run_adloc_ddp.py b/docs/run_adloc_ddp.py index aa9974f..aab2e38 100644 --- a/docs/run_adloc_ddp.py +++ b/docs/run_adloc_ddp.py @@ -285,17 +285,17 @@ for i in range(MAX_SST_ITER): if ddp_local_rank == 0: with torch.inference_mode(): - picks["residual_s"] = ( - travel_time(idx_sta, idx_eve, phase_type, phase_time, phase_weight)["residual_s"].detach().numpy() + picks["residual"] = ( + travel_time(idx_sta, idx_eve, phase_type, phase_time, phase_weight)["residual"].detach().numpy() ) - # station_term = picks.groupby("idx_sta").agg({"residual_s": "mean"}).reset_index() + # station_term = picks.groupby("idx_sta").agg({"residual": "mean"}).reset_index() station_term = ( picks.groupby("idx_sta") - .apply(lambda x: weighted_mean(x["residual_s"], x["phase_score"])) - .reset_index(name="residual_s") + .apply(lambda x: weighted_mean(x["residual"], x["phase_score"])) + .reset_index(name="residual") ) stations["station_term"] += ( - stations["idx_sta"].map(station_term.set_index("idx_sta")["residual_s"]).fillna(0) + stations["idx_sta"].map(station_term.set_index("idx_sta")["residual"]).fillna(0) ) raw_travel_time.station_dt.weight.data = torch.tensor( stations["station_term"].values, dtype=torch.float32 diff --git a/docs/run_ransac.py b/docs/run_ransac.py index 78d6b2b..a9e6f04 100644 --- a/docs/run_ransac.py +++ b/docs/run_ransac.py @@ -287,16 +287,16 @@ ) ## Separate P and S station term # station_term = ( - # picks[picks["mask"] == 1.0].groupby(["idx_sta", "phase_type"]).agg({"residual_s": "mean"}).reset_index() + # picks[picks["mask"] == 1.0].groupby(["idx_sta", "phase_type"]).agg({"residual_time": "mean"}).reset_index() # ) # stations["station_term_p"] = ( # stations["idx_sta"] - # .map(station_term[station_term["phase_type"] == 0].set_index("idx_sta")["residual_s"]) + # .map(station_term[station_term["phase_type"] == 0].set_index("idx_sta")["residual_time"]) # .fillna(0) # ) # stations["station_term_s"] = ( # stations["idx_sta"] - # .map(station_term[station_term["phase_type"] == 1].set_index("idx_sta")["residual_s"]) + # .map(station_term[station_term["phase_type"] == 1].set_index("idx_sta")["residual_time"]) # .fillna(0) # ) @@ -324,7 +324,7 @@ events.drop(["idx_eve", "x_km", "y_km", "z_km"], axis=1, inplace=True, errors="ignore") events.sort_values(["time"], inplace=True) - picks.rename({"mask": "adloc_mask", "residual_s": "adloc_residual_s"}, axis=1, inplace=True) + picks.rename({"mask": "adloc_mask", "residual_time": "adloc_residual_time", "residual_amplitude": "adloc_residual_amplitude"}, axis=1, inplace=True) picks["phase_type"] = picks["phase_type"].map({0: "P", 1: "S"}) picks.drop(["idx_eve", "idx_sta"], axis=1, inplace=True, errors="ignore") picks.sort_values(["phase_time"], inplace=True)