fix bug in the stanford test

AI4EPS · Jan 8, 2025 · 1cf129a · 1cf129a
1 parent c563845
commit 1cf129a
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 24 deletions.
diff --git a/adloc/adloc.py b/adloc/adloc.py
@@ -185,8 +185,8 @@ def calc_time(self, event_loc, station_loc, phase_type):
             z = event_loc[:, 2] - station_loc[:, 2]
             r = torch.sqrt(x**2 + y**2)
 
-            timetable = self.eikonal["up"] if phase_type == 0 else self.eikonal["us"]
-            timetable_grad = self.eikonal["grad_up"] if phase_type == 0 else self.eikonal["grad_us"]
+            timetable = self.eikonal["up"] if phase_type in [0, "P"] else self.eikonal["us"]
+            timetable_grad = self.eikonal["grad_up"] if phase_type in [0, "P"] else self.eikonal["grad_us"]
             timetable_grad_r = timetable_grad[0]
             timetable_grad_z = timetable_grad[1]
             rgrid0 = self.eikonal["rgrid"][0]
@@ -294,6 +294,7 @@ def __init__(
         # self.event_time.weight.requires_grad = True
 
     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)
             tt = dist / self.velocity[phase_type]
@@ -313,25 +314,25 @@ def calc_time(self, event_loc, station_loc, phase_type):
             z = event_loc[:, 2] - station_loc[:, 2]
             r = torch.sqrt(x**2 + y**2)
 
-            # timetable = self.eikonal["up"] if phase_type == 0 else self.eikonal["us"]
-            # timetable_grad = self.eikonal["grad_up"] if phase_type == 0 else self.eikonal["grad_us"]
-            # timetable_grad_r = timetable_grad[0]
-            # timetable_grad_z = timetable_grad[1]
-            # rgrid0 = self.eikonal["rgrid"][0]
-            # zgrid0 = self.eikonal["zgrid"][0]
-            # nr = self.eikonal["nr"]
-            # nz = self.eikonal["nz"]
-            # h = self.eikonal["h"]
-            # tt = CalcTravelTime.apply(r, z, timetable, timetable_grad_r, timetable_grad_z, rgrid0, zgrid0, nr, nz, h)
-
-            if phase_type in [0, "P"]:
-                timetable = self.timetable_p
-            elif phase_type in [1, "S"]:
-                timetable = self.timetable_s
-            else:
-                raise ValueError("phase_type should be 0 or 1. for P and S, respectively.")
-
-            tt = interp2d(timetable, r, z, self.rgrid, self.zgrid, self.h)
+            timetable = self.eikonal["up"] if phase_type in [0, "P"] else self.eikonal["us"]
+            timetable_grad = self.eikonal["grad_up"] if phase_type in [0, "P"] else self.eikonal["grad_us"]
+            timetable_grad_r = timetable_grad[0]
+            timetable_grad_z = timetable_grad[1]
+            rgrid0 = self.eikonal["rgrid"][0]
+            zgrid0 = self.eikonal["zgrid"][0]
+            nr = self.eikonal["nr"]
+            nz = self.eikonal["nz"]
+            h = self.eikonal["h"]
+            tt = CalcTravelTime.apply(r, z, timetable, timetable_grad_r, timetable_grad_z, rgrid0, zgrid0, nr, nz, h)
+
+            # if phase_type in [0, "P"]:
+            #     timetable = self.timetable_p
+            # elif phase_type in [1, "S"]:
+            #     timetable = self.timetable_s
+            # else:
+            #     raise ValueError("phase_type should be 0 or 1. for P and S, respectively.")
+
+            # tt = interp2d(timetable, r, z, self.rgrid, self.zgrid, self.h)
 
             tt = tt.float().view(nb1, ne1, 1)
 

diff --git a/docs/run_adloc.py b/docs/run_adloc.py
@@ -138,6 +138,7 @@
     # vp = [4.746, 4.793, 4.799, 5.045, 5.721, 5.879, 6.504, 6.708, 6.725, 7.800]
     # vs = [2.469, 2.470, 2.929, 2.930, 3.402, 3.403, 3.848, 3.907, 3.963, 4.500]
     # h = 0.3
+
     vel = {"Z": zz, "P": vp, "S": vs}
     config["eikonal"] = {
         "vel": vel,
@@ -200,7 +201,7 @@
         num_station,
         station_loc,
         event_loc=events[["x_km", "y_km", "z_km"]].values,
-        velocity={"P": vp, "S": vs},
+        velocity={"P": 6.0, "S": 6.0 / 1.73},
         eikonal=config["eikonal"],
     )
 
@@ -431,8 +432,11 @@
         tt = np.reshape(tt, (num_grid, num_picks))
         dt = tt - picks_per_event["travel_time"].values
         dt_mean = np.sum(dt * phase_weight, axis=-1, keepdims=True) / np.sum(phase_weight)
-        dt_std = np.sqrt(np.sum((dt - dt_mean) ** 2 * phase_weight, axis=-1) / np.sum(phase_weight))
-        idx = np.argmin(dt_std)
+        dt_std = np.sqrt(
+            np.sum((dt - dt_mean) ** 2 * phase_weight, axis=-1) / np.sum(phase_weight)
+        )  ## L2 loss/Huber loss
+        dt_abs = np.sum(np.abs(dt - dt_mean) * phase_weight, axis=-1) / np.sum(phase_weight)  ## L1 loss/Huber loss
+        idx = np.argmin(dt_abs)
         event_loc_gs.append(search_grid[idx] + event_loc0)
         event_time_gs.append(np.mean(dt, axis=-1)[idx] + event_time0)
 
@@ -446,6 +450,8 @@
         uncertainty = np.sqrt(variance)
         event_uncertainty.append(uncertainty)
 
+        # import matplotlib.pyplot as plt
+
         # tt = np.reshape(tt, (nx, ny, nz, num_picks))
         # dt = np.reshape(dt, (nx, ny, nz, num_picks))
         # dt = np.sum(np.abs(dt) * phase_weight, axis=-1) / np.sum(phase_weight)
@@ -470,6 +476,7 @@
         # ax[2].invert_yaxis()
         # plt.savefig("debug.png", bbox_inches="tight", dpi=300)
         # plt.close(fig)
+        # raise
 
     events = events_init.copy()
     events["time"] = events["time"] + pd.to_timedelta(np.squeeze(event_time_gs), unit="s")