stanford test works

adloc/adloc.py

@@ -25,12 +25,15 @@ def interp2d(time_table, x, y, xgrid, ygrid, h):
     ny = len(ygrid)
     assert time_table.shape == (nx, ny)
 
-    ix0 = torch.floor((x - xgrid[0]) / h).clamp(0, nx - 2).long()
-    iy0 = torch.floor((y - ygrid[0]) / h).clamp(0, ny - 2).long()
-    ix1 = ix0 + 1
-    iy1 = iy0 + 1
-    # x = (torch.clamp(x, self.xgrid[0], self.xgrid[-1]) - self.xgrid[0]) / self.h
-    # y = (torch.clamp(y, self.ygrid[0], self.ygrid[-1]) - self.ygrid[0]) / self.h
+    with torch.no_grad():
+        ix0 = torch.floor((x - xgrid[0]) / h).clamp(0, nx - 2).long()
+        iy0 = torch.floor((y - ygrid[0]) / h).clamp(0, ny - 2).long()
+        ix1 = ix0 + 1
+        iy1 = iy0 + 1
+
+    # x = (torch.clamp(x, xgrid[0], xgrid[-1]) - xgrid[0]) / h
+    # y = (torch.clamp(y, ygrid[0], ygrid[-1]) - ygrid[0]) / h
+
     x = (clamp(x, xgrid[0], xgrid[-1]) - xgrid[0]) / h
     y = (clamp(y, ygrid[0], ygrid[-1]) - ygrid[0]) / h
 
@@ -220,6 +223,9 @@ def __init__(
             grad_type=grad_type,
         )
 
+        self.event_loc.weight.requires_grad = True
+        self.event_time.weight.requires_grad = False
+
     def calc_time(self, event_loc, station_loc, phase_type):
         if self.eikonal is None:
             dist = torch.linalg.norm(event_loc - station_loc, axis=-1, keepdim=True)

adloc/data.py

@@ -161,7 +161,8 @@ def read_data(self):
                 if key == -1:
                     continue
                 group = self.picks_by_event.get_group(key)
-                phase_time.append(group["travel_time"].values)
+                phase_time.append(group["travel_time"].values)  # seconds
+                # phase_time.append(group["phase_time"].values) # datetime not supported
                 phase_score.append(group["phase_score"].values)
                 phase_type.extend(group["phase_type"].values.tolist())
                 idx_eve.extend(group["idx_eve"].values.tolist())

adloc/eikonal2d.py

@@ -179,7 +179,7 @@ def traveltime(event_index, station_index, phase_type, events, stations, eikonal
 
         tt = np.zeros(len(phase_type), dtype=np.float32)
 
-        if isinstance(phase_type[0].item(), str):
+        if isinstance(phase_type[0], str):
             p_index = phase_type == "P"
             s_index = phase_type == "S"
         elif isinstance(phase_type[0].item(), int):
@@ -293,12 +293,26 @@ def init_eikonal2d(config):
     nz = len(zgrid)
 
     vel = config["vel"]
-    zz, vp, vs = vel["Z"], vel["P"], vel["S"]
+    zz, vp, vs = np.array(vel["Z"]), np.array(vel["P"]), np.array(vel["S"])
+
+    # ##############################################
+    # ## make the velocity staircase not linear 
+    # zz_grid = zz[1:] - h
+    # vp_grid = vp[:-1]
+    # vs_grid = vs[:-1]
+    # zz = np.concatenate([zz, zz_grid])
+    # vp = np.concatenate([vp, vp_grid])
+    # vs = np.concatenate([vs, vs_grid])
+    # idx = np.argsort(zz)
+    # zz = zz[idx]
+    # vp = vp[idx]
+    # vs = vs[idx]
+    # ##############################################
+
     vp1d = np.interp(zgrid, zz, vp)
     vs1d = np.interp(zgrid, zz, vs)
     vp = np.tile(vp1d, (nr, 1))
     vs = np.tile(vs1d, (nr, 1))
-
     # ir0 = np.floor(config["source_loc"][0] / h).astype(np.int64)
     # iz0 = np.floor(config["source_loc"][1] / h).astype(np.int64)
     ir0 = np.round(0 - rlim[0] / h).astype(np.int64)

adloc/sacloc2d.py

@@ -82,16 +82,14 @@ def huber_loss_grad(event, X, y, vel={0: 6.0, 1: 6.0 / 1.75}, stations=None, eik
 
         station_index = X[:, 0].astype(int)
         phase_type = X[:, 1].astype(int)
-        if X.shape[1] > 2:
-            phase_weight = X[:, 2]
-        else:
-            phase_weight = np.ones(len(X))
+        phase_weight = X[:, 2]
 
         if eikonal is None:
             v = np.array([vel[t] for t in phase_type])
             tt = np.linalg.norm(event[:3] - stations[station_index, :3], axis=-1) / v + event[3]
         else:
             tt = traveltime(0, station_index, phase_type, event[np.newaxis, :3], stations, eikonal) + event[3]
+
         t_diff = tt - y
         l1 = np.squeeze((np.abs(t_diff) > sigma))
         l2 = np.squeeze((np.abs(t_diff) <= sigma))
@@ -123,7 +121,7 @@ def fit(self, X, y=None, event_index=0):
             yt = y[:, 0]  # time
             ya = y[:, 1]  # amplitude
         else:
-            yt = y
+            yt = y[:, 0]
 
         opt = scipy.optimize.minimize(
             # self.huber_loss_grad,
@@ -146,10 +144,7 @@ def fit(self, X, y=None, event_index=0):
 
         if self.use_amplitude:
             station_index = X[:, 0].astype(int)
-            if X.shape[1] > 2:
-                phase_weight = X[:, 2:3]
-            else:
-                phase_weight = np.ones(len(X))
+            phase_weight = X[:, 2:3]
             self.magnitudes[event_index] = calc_mag(
                 ya[:, np.newaxis], self.events[event_index, :3], self.stations[station_index, :3], phase_weight
             )
@@ -181,7 +176,7 @@ def predict(self, X, event_index=0):
             ).squeeze()
             return np.array([tt, amp]).T
         else:
-            return tt
+            return np.array([tt]).T
 
     def score(self, X, y=None, event_index=0):
         """

adloc/utils.py

@@ -31,13 +31,22 @@ def is_data_valid(X, y):
 
     MIN_PICKS = config["min_picks"]
     MIN_PICKS_RATIO = config["min_picks_ratio"]
-    MAX_RESIDUAL = [config["max_residual_time"], config["max_residual_amplitude"]]
+    if config["use_amplitude"]:
+        MAX_RESIDUAL = [config["max_residual_time"], config["max_residual_amplitude"]]
+    else:
+        MAX_RESIDUAL = config["max_residual_time"]
 
-    X = picks.merge(
-        stations[["x_km", "y_km", "z_km", "station_id", "station_term_time", "station_term_amplitude"]],
-        # stations[["x_km", "y_km", "z_km", "station_id", "station_term_p", "station_term_s"]], ## Separate P and S station term
-        on="station_id",
-    )
+    if config["use_amplitude"]:
+        X = picks.merge(
+            stations[["x_km", "y_km", "z_km", "station_id", "station_term_time", "station_term_amplitude"]],
+            # stations[["x_km", "y_km", "z_km", "station_id", "station_term_p", "station_term_s"]], ## Separate P and S station term
+            on="station_id",
+        )
+    else:
+        X = picks.merge(
+            stations[["x_km", "y_km", "z_km", "station_id", "station_term_time"]],
+            on="station_id",
+        )
     t0 = X["phase_time"].min()
 
     if event_init is not None:
@@ -49,13 +58,12 @@ def is_data_valid(X, y):
     # ystd = np.std(X["y_km"])
     # rstd = np.sqrt(xstd**2 + ystd**2)
 
-    X.rename(
-        columns={"phase_type": "type", "phase_score": "score", "phase_time": "t_s", "phase_amplitude": "amp"},
-        inplace=True,
-    )
+    X.rename(columns={"phase_type": "type", "phase_score": "score", "phase_time": "t_s"}, inplace=True)
     X["t_s"] = (X["t_s"] - t0).dt.total_seconds()
     X["t_s"] = X["t_s"] - X["station_term_time"]
-    X["amp"] = X["amp"] - X["station_term_amplitude"]
+    if config["use_amplitude"]:
+        X.rename(columns={"phase_amplitude": "amp"}, inplace=True)
+        X["amp"] = X["amp"] - X["station_term_amplitude"]
     # X["t_s"] = X.apply(
     #     lambda x: x["t_s"] - x["station_term_p"] if x["type"] == 0 else x["t_s"] - x["station_term_s"], axis=1
     # ) ## Separate P and S station term
@@ -80,15 +88,19 @@ def is_data_valid(X, y):
         is_data_valid=is_data_valid,
     )
     try:
-        reg.fit(X[["idx_sta", "type", "score", "amp"]].values, X[["t_s", "amp"]].values)
+        if config["use_amplitude"]:
+            reg.fit(X[["idx_sta", "type", "score", "amp"]].values, X[["t_s", "amp"]].values)
+        else:
+            reg.fit(X[["idx_sta", "type", "score"]].values, X[["t_s"]].values)
     except Exception as e:
         print(f"No valid model for event_index {event_index}.")
         message = "RANSAC could not find a valid consensus set."
         if str(e)[: len(message)] != message:
             print(e)
         picks["mask"] = 0
         picks["residual_time"] = 0.0
-        picks["residual_amplitude"] = 0.0
+        if config["use_amplitude"]:
+            picks["residual_amplitude"] = 0.0
         return picks, None
 
     estimator = reg.estimator_
@@ -100,10 +112,13 @@ def is_data_valid(X, y):
         tt = output[:, 0]
         amp = output[:, 1]
     else:
-        tt = output
-    score = estimator.score(
-        X[["idx_sta", "type", "score", "amp"]].values[inlier_mask], y=X[["t_s", "amp"]].values[inlier_mask]
-    )
+        tt = output[:, 0]
+    if config["use_amplitude"]:
+        score = estimator.score(
+            X[["idx_sta", "type", "score", "amp"]].values[inlier_mask], y=X[["t_s", "amp"]].values[inlier_mask]
+        )
+    else:
+        score = estimator.score(X[["idx_sta", "type", "score"]].values[inlier_mask], y=X[["t_s"]].values[inlier_mask])
 
     if (np.sum(inlier_mask) > MIN_PICKS) and (score > config["min_score"]):
         mean_residual_time = np.sum(np.abs(X["t_s"].values - tt) * inlier_mask) / np.sum(inlier_mask)
@@ -121,9 +136,10 @@ def is_data_valid(X, y):
             "time": t0 + pd.Timedelta(t, unit="s"),
             "adloc_score": score,
             "adloc_residual_time": mean_residual_time,
-            "adloc_residual_amplitude": mean_residual_amp,
             "num_picks": np.sum(inlier_mask),
         }
+        if config["use_amplitude"]:
+            event["adloc_residual_amplitude"] = mean_residual_amp
         event = pd.DataFrame([event])
     else:
         inlier_mask = np.zeros(len(inlier_mask), dtype=bool)
@@ -182,12 +198,13 @@ def invert_location(picks, stations, config, estimator, events_init=None, iter=0
         for idx_eve, picks_by_event in picks.groupby("idx_eve"):
             if events_init is not None:
                 event_init = events_init[events_init["idx_eve"] == idx_eve]
-                if len(event_init) > 0:
+                if len(event_init) == 1:
                     event_init = event_init[["x_km", "y_km", "z_km", "time"]].values[0]
+                elif len(event_init) > 1:
+                    event_init = event_init[["x_km", "y_km", "z_km", "time"]].values[0]
+                    print(f"Multiple initial locations for event_index {idx_eve}.")
                 else:
                     event_init = None
-            else:
-                event_init = None
 
             # picks_, event_ = invert(picks_by_event, stations, config, estimator, idx_eve, event_init)
             # if event_ is not None:
@@ -215,8 +232,6 @@ def invert_location(picks, stations, config, estimator, events_init=None, iter=0
     picks_inverted = pd.concat(picks_inverted)
     if events_init is not None:
         events_inverted = events_inverted.merge(events_init[["event_index", "idx_eve"]], on="idx_eve")
-    else:
-        events_inverted["event_index"] = events_inverted["idx_eve"]
 
     print(f"ADLoc using {len(picks_inverted[picks_inverted['mask'] == 1])} picks outof {len(picks_inverted)} picks")