refactor intensity prediction functions and update Gradio interface for epicenter input

app.py

@@ -405,28 +405,166 @@ def calculate_intensity(pga, label=False):
 # ============ Gradio 介面函數 ============
 
 def load_waveform(event_name):
+    """載入完整的 mseed 檔案（包含所有測站）"""
     file_path = EARTHQUAKE_EVENTS[event_name]
     st = read(file_path)
-    tr = st[0]
-    times = tr.times()
-    data = tr.data
-    return times, data, tr.stats.sampling_rate
+    return st
 
 
-def plot_waveform(times, data, start_time, end_time, sampling_rate):
-    fig, ax = plt.subplots(figsize=(12, 3))
-    ax.plot(times, data, 'gray', linewidth=0.5, alpha=0.6)
+def calculate_distance(lat1, lon1, lat2, lon2):
+    """計算兩點間的距離（簡化的平面距離，單位：度）"""
+    return np.sqrt((lat1 - lat2)**2 + (lon1 - lon2)**2)
 
-    mask = (times >= start_time) & (times <= end_time)
-    ax.plot(times[mask], data[mask], 'blue', linewidth=1)
 
-    ax.axvline(start_time, color='red', linestyle='--', linewidth=1)
-    ax.axvline(end_time, color='red', linestyle='--', linewidth=1)
+def select_nearest_stations(st, epicenter_lat, epicenter_lon, n_stations=25):
+    """選擇距離震央最近的 n 個測站"""
+    station_distances = []
 
-    ax.set_xlabel('Time (s)')
-    ax.set_ylabel('Amplitude')
-    ax.set_title('Seismic Waveform')
-    ax.grid(True, alpha=0.3)
+    # 計算每個測站到震央的距離
+    for tr in st:
+        station_code = tr.stats.station
+
+        # 從 site_info 中查詢測站位置
+        try:
+            station_data = site_info[site_info["Station"] == station_code]
+            if len(station_data) == 0:
+                continue
+
+            lat = station_data["Latitude"].values[0]
+            lon = station_data["Longitude"].values[0]
+            elev = station_data["Elevation"].values[0]
+
+            distance = calculate_distance(epicenter_lat, epicenter_lon, lat, lon)
+            station_distances.append({
+                "station": station_code,
+                "distance": distance,
+                "latitude": lat,
+                "longitude": lon,
+                "elevation": elev
+            })
+        except Exception as e:
+            logger.warning(f"測站 {station_code} 資訊查詢失敗: {e}")
+            continue
+
+    # 按距離排序並選擇最近的 n 個
+    station_distances.sort(key=lambda x: x["distance"])
+    selected_stations = station_distances[:n_stations]
+
+    logger.info(f"選擇了 {len(selected_stations)} 個最近的測站")
+    return selected_stations
+
+
+def extract_waveforms_from_stream(st, selected_stations, start_time, end_time, vs30_input):
+    """從 Stream 中提取選定測站的波形資料"""
+    waveforms = []
+    station_info_list = []
+    valid_stations = []
+
+    sampling_rate = 100  # 假設 100 Hz
+    start_idx = int(start_time * sampling_rate)
+    end_idx = int(end_time * sampling_rate)
+    target_length = 3000
+
+    for station_data in selected_stations:
+        station_code = station_data["station"]
+
+        try:
+            # 選擇該測站的所有分量
+            st_station = st.select(station=station_code)
+
+            if len(st_station) == 0:
+                continue
+
+            # 嘗試取得 Z, N, E 分量
+            z_trace = st_station.select(component="Z")
+            n_trace = st_station.select(component="N") or st_station.select(component="1")
+            e_trace = st_station.select(component="E") or st_station.select(component="2")
+
+            # 如果沒有三分量，使用 Z 分量重複
+            if len(z_trace) > 0:
+                z_data = z_trace[0].data[start_idx:end_idx]
+            else:
+                continue
+
+            if len(n_trace) > 0:
+                n_data = n_trace[0].data[start_idx:end_idx]
+            else:
+                n_data = z_data.copy()
+
+            if len(e_trace) > 0:
+                e_data = e_trace[0].data[start_idx:end_idx]
+            else:
+                e_data = z_data.copy()
+
+            # 訊號處理
+            z_data = signal_processing(z_data)
+            n_data = signal_processing(n_data)
+            e_data = signal_processing(e_data)
+
+            # 調整長度到 3000
+            for data in [z_data, n_data, e_data]:
+                if len(data) > target_length:
+                    data = data[:target_length]
+                elif len(data) < target_length:
+                    data = np.pad(data, (0, target_length - len(data)))
+
+            # 組合三分量 (3000, 3)
+            waveform_3c = np.stack([z_data[:target_length],
+                                   n_data[:target_length],
+                                   e_data[:target_length]], axis=1)
+            waveforms.append(waveform_3c)
+
+            # 準備測站資訊
+            vs30 = get_vs30(station_data["latitude"], station_data["longitude"], vs30_input)
+            station_info_list.append([
+                station_data["latitude"],
+                station_data["longitude"],
+                station_data["elevation"],
+                vs30
+            ])
+            valid_stations.append(station_data)
+
+        except Exception as e:
+            logger.warning(f"測站 {station_code} 波形提取失敗: {e}")
+            continue
+
+    logger.info(f"成功提取 {len(waveforms)} 個測站的波形")
+    return waveforms, station_info_list, valid_stations
+
+
+def plot_waveform(st, selected_stations, start_time, end_time):
+    """繪製選定測站的波形圖（堆疊顯示前 10 個測站）"""
+    fig, axes = plt.subplots(min(10, len(selected_stations)), 1, figsize=(12, 8), sharex=True)
+
+    if len(selected_stations) == 1:
+        axes = [axes]
+
+    for i, station_data in enumerate(selected_stations[:10]):
+        station_code = station_data["station"]
+
+        try:
+            st_station = st.select(station=station_code)
+            if len(st_station) > 0:
+                tr = st_station[0]
+                times = tr.times()
+                data = tr.data
+
+                axes[i].plot(times, data, 'black', linewidth=0.5)
+
+                # 標記選取範圍
+                axes[i].axvline(start_time, color='red', linestyle='--', linewidth=1, alpha=0.7)
+                axes[i].axvline(end_time, color='red', linestyle='--', linewidth=1, alpha=0.7)
+                axes[i].axvspan(start_time, end_time, alpha=0.2, color='blue')
+
+                axes[i].set_ylabel(f'{station_code}\n({station_data["distance"]:.2f}°)', fontsize=8)
+                axes[i].grid(True, alpha=0.3)
+                axes[i].tick_params(labelsize=8)
+        except Exception as e:
+            logger.warning(f"無法繪製測站 {station_code}: {e}")
+
+    axes[-1].set_xlabel('Time (s)')
+    fig.suptitle(f'波形記錄（前 10 個最近測站，共選擇 {len(selected_stations)} 個）', fontsize=12)
+    plt.tight_layout()
 
     return fig
 
@@ -463,61 +601,83 @@ def plot_intensity_map(pga_list, target_names):
     return fig
 
 
-def predict_intensity(event_name, start_time, end_time, lon, lat, vs30_input):
-    # 1. 載入波形
-    times, data, sampling_rate = load_waveform(event_name)
+def predict_intensity(event_name, start_time, end_time, epicenter_lon, epicenter_lat, vs30_input):
+    """執行震度預測"""
+    try:
+        # 1. 載入完整的 mseed 檔案
+        logger.info(f"載入地震事件: {event_name}")
+        st = load_waveform(event_name)
+        logger.info(f"載入了 {len(st)} 個 trace")
 
-    # 2. 切片波形
-    start_idx = int(start_time * sampling_rate)
-    end_idx = int(end_time * sampling_rate)
-    waveform_slice = data[start_idx:end_idx]
-
-    # 3. 訊號處理
-    waveform_processed = signal_processing(waveform_slice)
-
-    # 4. 準備模型輸入
-    # 假設單測站三軸資料（這裡簡化為重複使用Z軸）
-    waveform_3c = np.array(
-        [[waveform_processed, waveform_processed, waveform_processed]])
-    waveform_3c = waveform_3c.transpose(0, 2, 1)  # (1, 3000, 3)
-
-    # 準備測站資訊（使用使用者輸入的 Vs30 或資料庫值）
-    vs30 = get_vs30(lat, lon, vs30_input)
-    station_info_input = np.array([[lat, lon, 100, vs30]])  # elevation 假設 100m
-
-    # 準備目標測站資訊
-    target_list = []
-    target_names = []
-    for target in target_dict[:25]:  # 限制25個目標
-        target_list.append([target["latitude"], target["longitude"],
-                            target["elevation"],
-                            get_vs30(target["latitude"], target["longitude"], vs30_input)])
-        target_names.append(target["station"])
-
-    # 組合成 tensor
-    tensor_data = {
-        "waveform": torch.tensor(waveform_3c).unsqueeze(0).double(),
-        "station": torch.tensor(station_info_input).unsqueeze(0).double(),
-        "target": torch.tensor(target_list).unsqueeze(0).double(),
-    }
-
-    # 5. 執行預測
-    with torch.no_grad():
-        weight, sigma, mu = model(tensor_data)
-        pga_list = torch.sum(weight * mu,
-                             dim=2).cpu().detach().numpy().flatten().tolist()
-
-    # 6. 繪製結果
-    waveform_plot = plot_waveform(times, data, start_time, end_time, sampling_rate)
-    intensity_plot = plot_intensity_map(pga_list, target_names)
-
-    # 統計資訊
-    max_intensity = max([calculate_intensity(pga, label=True) for pga in pga_list])
-    stats = f"選取時間範圍: {start_time:.1f} - {end_time:.1f} 秒\n"
-    stats += f"測站位置: ({lon:.4f}, {lat:.4f})\n"
-    stats += f"預測最大震度: {max_intensity}"
-
-    return waveform_plot, intensity_plot, stats
+        # 2. 根據震央距離選擇最近的 25 個測站
+        logger.info(f"選擇距離震央 ({epicenter_lat}, {epicenter_lon}) 最近的測站...")
+        selected_stations = select_nearest_stations(st, epicenter_lat, epicenter_lon, n_stations=25)
+
+        if len(selected_stations) == 0:
+            return None, None, "錯誤：找不到有效的測站資料"
+
+        # 3. 從選定的測站提取波形
+        logger.info(f"提取波形資料（時間範圍: {start_time}-{end_time} 秒）...")
+        waveforms, station_info_list, valid_stations = extract_waveforms_from_stream(
+            st, selected_stations, start_time, end_time, vs30_input
+        )
+
+        if len(waveforms) == 0:
+            return None, None, "錯誤：無法提取波形資料"
+
+        # 4. Padding 到 25 個測站（模型要求）
+        max_stations = 25
+        waveform_padded = np.zeros((max_stations, 3000, 3))
+        station_info_padded = np.zeros((max_stations, 4))
+
+        for i in range(min(len(waveforms), max_stations)):
+            waveform_padded[i] = waveforms[i]
+            station_info_padded[i] = station_info_list[i]
+
+        # 5. 準備目標測站資訊
+        target_list = []
+        target_names = []
+        for target in target_dict[:25]:
+            target_list.append([
+                target["latitude"],
+                target["longitude"],
+                target["elevation"],
+                get_vs30(target["latitude"], target["longitude"], vs30_input)
+            ])
+            target_names.append(target["station"])
+
+        # 6. 組合成 tensor
+        tensor_data = {
+            "waveform": torch.tensor(waveform_padded).unsqueeze(0).double(),
+            "station": torch.tensor(station_info_padded).unsqueeze(0).double(),
+            "target": torch.tensor(target_list).unsqueeze(0).double(),
+        }
+
+        # 7. 執行預測
+        logger.info("執行模型預測...")
+        with torch.no_grad():
+            weight, sigma, mu = model(tensor_data)
+            pga_list = torch.sum(weight * mu, dim=2).cpu().detach().numpy().flatten().tolist()
+
+        # 8. 繪製結果
+        waveform_plot = plot_waveform(st, selected_stations, start_time, end_time)
+        intensity_plot = plot_intensity_map(pga_list, target_names)
+
+        # 9. 統計資訊
+        max_intensity = max([calculate_intensity(pga, label=True) for pga in pga_list])
+        stats = f"選取時間範圍: {start_time:.1f} - {end_time:.1f} 秒\n"
+        stats += f"震央位置: ({epicenter_lon:.4f}, {epicenter_lat:.4f})\n"
+        stats += f"使用測站數: {len(waveforms)} / 25\n"
+        stats += f"預測最大震度: {max_intensity}"
+
+        logger.info("預測完成！")
+        return waveform_plot, intensity_plot, stats
+
+    except Exception as e:
+        logger.error(f"預測過程發生錯誤: {e}")
+        import traceback
+        traceback.print_exc()
+        return None, None, f"錯誤: {str(e)}"
 
 
 # ============ Gradio 介面 ============
@@ -539,11 +699,12 @@ with gr.Blocks(title="TTSAM 震度預測系統") as demo:
                 start_slider = gr.Slider(0, 300, value=0, step=1, label="起始時間 (秒)")
                 end_slider = gr.Slider(0, 300, value=30, step=1, label="結束時間 (秒)")
 
-            gr.Markdown("### 測站位置")
+            gr.Markdown("### 震央位置")
             with gr.Row():
-                lon_input = gr.Number(value=121.5, label="經度")
-                lat_input = gr.Number(value=24.0, label="緯度")
+                epicenter_lon_input = gr.Number(value=121.5, label="震央經度")
+                epicenter_lat_input = gr.Number(value=24.0, label="震央緯度")
 
+            gr.Markdown("### 場址參數")
             with gr.Row():
                 vs30_input = gr.Number(
                     value=600,
@@ -553,6 +714,13 @@ with gr.Blocks(title="TTSAM 震度預測系統") as demo:
 
             predict_btn = gr.Button("🔮 執行預測", variant="primary")
 
+            gr.Markdown("""
+            ### 說明
+            - 系統會根據震央位置自動選擇最近的 25 個測站
+            - 從選定的時間範圍提取波形資料（30 秒）
+            - 預測全台灣目標測站的震度分布
+            """)
+
         # 右側：震度分布圖
         with gr.Column(scale=1):
             gr.Markdown("## 預測震度分布")
@@ -568,7 +736,7 @@ with gr.Blocks(title="TTSAM 震度預測系統") as demo:
     # 綁定事件
     predict_btn.click(
         fn=predict_intensity,
-        inputs=[event_dropdown, start_slider, end_slider, lon_input, lat_input, vs30_input],
+        inputs=[event_dropdown, start_slider, end_slider, epicenter_lon_input, epicenter_lat_input, vs30_input],
         outputs=[waveform_plot, intensity_plot, stats_output]
     )