from pathlib import Path import numpy as np import pandas as pd import torch from torch.utils.data import Dataset # --------------------------------------------------------------------- # 1. Sparse Dataset 정의 (panel + static → ST-ResNet 입력 생성) # --------------------------------------------------------------------- class GridSTSparseDataset(Dataset): """ ST-ResNet 계열 모델용 시공간 Grid 데이터셋. 동적(dynamic) 피처 (F_dyn = 3): 1) cnt(g,t) : 사고 건수 2) dist(g,t) : 평균 현장거리 (mean_dist, meter) 3) traffic(g,t): GRID_SPEED (평균 교통량/속도) → 한 시점 동적 텐서 shape: (H, W, 3) 예) closeness 입력: (L_c, H, W, 3) 정적(static) 피처 (F_stat = 3): 1) road_len_in_grid(g) 2) n_links(g) 3) n_segments(g) → 정적 텐서 shape: (H, W, 3) 반환 dict: { "X_closeness": (L_c, H, W, 3), "X_period": (L_p, H, W, 3) or None, "X_trend": (L_t, H, W, 3) or None, "X_static": (H, W, 3), "y": (H, W), # 평균 소요시간 "mask": (H, W), # 관측 마스크 "t": int, # t_idx "ext": (3,) or None, # [hour/23, dow/6, is_weekend] } """ def __init__( self, panel: pd.DataFrame, static: np.ndarray, T: int, R: int, C: int, seq_len: int = 6, period: int = 24, trend: int = 24 * 7, num_periods: int = 3, num_trends: int = 2, hour_arr: np.ndarray | None = None, dow_arr: np.ndarray | None = None, is_weekend_arr: np.ndarray | None = None, ): """ panel 컬럼 요구사항: - "t_idx" - "r_idx", "c_idx" - "cnt" : 사고 건수 - "mean_time" : 평균 소요시간 (초) - "mean_dist" : 평균 현장거리 - "traffic" : GRID_SPEED (평균 속도/교통량) """ super().__init__() self.panel = panel self.static = torch.from_numpy(static).float() # (R, C, F_stat=3) self.T = T self.R = R self.C = C self.seq_len = seq_len self.period = period self.trend = trend self.num_periods = num_periods self.num_trends = num_trends self.hour_arr = hour_arr self.dow_arr = dow_arr self.is_weekend_arr = is_weekend_arr # (t_idx → sparse grid 정보) 캐시 # time_dict[t] = (r_idx, c_idx, cnt, mean_time, mean_dist, traffic) self._build_time_index() # temporal dependency 만족하는 t만 유효 min_t = seq_len if period > 0 and num_periods > 0: min_t = max(min_t, num_periods * period) if trend > 0 and num_trends > 0: min_t = max(min_t, num_trends * trend) self.min_t = min_t self.valid_t = list(range(self.min_t, self.T)) print( f"[GridSTSparseDataset] T={self.T}, R={self.R}, C={self.C}, " f"seq_len={self.seq_len}, period={self.period}, trend={self.trend}, " f"|valid_t|={len(self.valid_t)}" ) def _build_time_index(self): """ panel를 t_idx별로 그룹핑해서 sparse grid 저장. time_dict[t] = ( r_idx: int[], c_idx: int[], cnt: float32[], mean_time: float32[], mean_dist: float32[], traffic: float32[], ) """ self.time_dict = {} g = self.panel.groupby("t_idx") for t, df_t in g: r = df_t["r_idx"].to_numpy(dtype=int) c = df_t["c_idx"].to_numpy(dtype=int) cnt = df_t["cnt"].to_numpy(dtype=np.float32) mt = df_t["mean_time"].to_numpy(dtype=np.float32) dist = df_t["mean_dist"].to_numpy(dtype=np.float32) traf = df_t["traffic"].to_numpy(dtype=np.float32) self.time_dict[int(t)] = (r, c, cnt, mt, dist, traf) def __len__(self): return len(self.valid_t) def _build_map_at_t(self, t: int): """ 시점 t에 대해 full grid 맵 생성. Returns ------- cnt_map : (R, C), 사고 건수 dist_map : (R, C), 평균 현장거리 traf_map : (R, C), 평균 교통량/속도 y_map : (R, C), 평균 소요시간 (초) mask_map : (R, C), 관측 마스크 (0/1) """ R, C = self.R, self.C cnt_map = np.zeros((R, C), dtype=np.float32) dist_map = np.zeros((R, C), dtype=np.float32) traf_map = np.zeros((R, C), dtype=np.float32) y_map = np.zeros((R, C), dtype=np.float32) if t in self.time_dict: r_idx, c_idx, cnt, mt, dist, traf = self.time_dict[t] cnt_map[r_idx, c_idx] = cnt dist_map[r_idx, c_idx] = dist traf_map[r_idx, c_idx] = traf y_map[r_idx, c_idx] = mt # 0 < y < 3600 인 지점만 유효 관측 mask_map = ((y_map > 0) & (y_map < 3600)).astype(np.float32) return cnt_map, dist_map, traf_map, y_map, mask_map def __getitem__(self, idx): t = self.valid_t[idx] # ---------- closeness (L_c 시점, F_dyn=3) ---------- Xc_cnt = [] Xc_dist = [] Xc_traf = [] for k in range(self.seq_len, 0, -1): t_k = t - k cnt_map, dist_map, traf_map, _, _ = self._build_map_at_t(t_k) Xc_cnt.append(cnt_map) Xc_dist.append(dist_map) Xc_traf.append(traf_map) # (L_c, R, C, 3) = [cnt, dist, traffic] X_c = np.stack( [ np.stack(Xc_cnt, axis=0), np.stack(Xc_dist, axis=0), np.stack(Xc_traf, axis=0), ], axis=-1, ) # ---------- period ---------- if self.period > 0 and self.num_periods > 0: Xp_cnt = [] Xp_dist = [] Xp_traf = [] for k in range(1, self.num_periods + 1): t_k = t - k * self.period cnt_map, dist_map, traf_map, _, _ = self._build_map_at_t(t_k) Xp_cnt.append(cnt_map) Xp_dist.append(dist_map) Xp_traf.append(traf_map) # (L_p, R, C, 3) X_p = np.stack( [ np.stack(Xp_cnt, axis=0), np.stack(Xp_dist, axis=0), np.stack(Xp_traf, axis=0), ], axis=-1, ) else: X_p = None # ---------- trend ---------- if self.trend > 0 and self.num_trends > 0: Xt_cnt = [] Xt_dist = [] Xt_traf = [] for k in range(1, self.num_trends + 1): t_k = t - k * self.trend cnt_map, dist_map, traf_map, _, _ = self._build_map_at_t(t_k) Xt_cnt.append(cnt_map) Xt_dist.append(dist_map) Xt_traf.append(traf_map) # (L_t, R, C, 3) X_t = np.stack( [ np.stack(Xt_cnt, axis=0), np.stack(Xt_dist, axis=0), np.stack(Xt_traf, axis=0), ], axis=-1, ) else: X_t = None # ---------- target & mask ---------- _, _, _, y_map, mask_map = self._build_map_at_t(t) # ---------- static ---------- X_s = self.static # (R, C, 3): [road_len_in_grid, n_links, n_segments] # ---------- 외생 변수 ---------- if ( self.hour_arr is not None and self.dow_arr is not None and self.is_weekend_arr is not None ): h = float(self.hour_arr[t]) / 23.0 d = float(self.dow_arr[t]) / 6.0 w = float(self.is_weekend_arr[t]) ext = torch.tensor([h, d, w], dtype=torch.float32) else: ext = None # Torch 변환 X_c = torch.from_numpy(X_c).float() # (L_c, R, C, 3) y_t = torch.from_numpy(y_map).float() # (R, C) m_t = torch.from_numpy(mask_map).float() # (R, C) X_s = X_s.float() # (R, C, 3) if X_p is not None: X_p = torch.from_numpy(X_p).float() # (L_p, R, C, 3) if X_t is not None: X_t = torch.from_numpy(X_t).float() # (L_t, R, C, 3) return { "X_closeness": X_c, "X_period": X_p, "X_trend": X_t, "X_static": X_s, "y": y_t, "mask": m_t, "t": t, "ext": ext, } # --------------------------------------------------------------------- # 2. 해상도별 parquet → panel / static / 시간메타 생성 # --------------------------------------------------------------------- def build_sparse_panel_and_static(GRID_RES_M: int): """ 준비된 parquet에서 모델 입력용 panel / static / 시간 메타데이터 생성. 동적(panel) (t_idx, r_idx, c_idx): - cnt : cnt_events (사고 건수) - mean_time : 소요시각 평균 - mean_dist : 현장거리 평균 - traffic : GRID_SPEED (평균 교통량/속도) 정적(static) (R, C, 3): - road_len_in_grid - n_links - n_segments """ BASE_DIR = Path(__file__).resolve().parent DB_DIR = BASE_DIR / "DB" PREP_DIR = DB_DIR / "prepared" PATH_FACT_ACC_GRID = PREP_DIR / f"fact_accident_dispatch_grid_{GRID_RES_M}m.parquet" PATH_DYN_FEAT = PREP_DIR / f"dyn_features_grid_{GRID_RES_M}m.parquet" PATH_GRID_STATIC = PREP_DIR / f"grid_static_road_{GRID_RES_M}m.parquet" PATH_TIME_BIAS = PREP_DIR / "time_bias_features.parquet" print(f"[{GRID_RES_M}m] Loading prepared parquet files...") fact = pd.read_parquet(PATH_FACT_ACC_GRID) dyn = pd.read_parquet(PATH_DYN_FEAT) grid_static = pd.read_parquet(PATH_GRID_STATIC) time_bias = pd.read_parquet(PATH_TIME_BIAS) print(f"[{GRID_RES_M}m] fact shape:", fact.shape) print(f"[{GRID_RES_M}m] dyn shape:", dyn.shape) print(f"[{GRID_RES_M}m] grid_static shape:", grid_static.shape) print(f"[{GRID_RES_M}m] time_bias shape:", time_bias.shape) # r_idx / c_idx 준비 if "r_idx" in fact.columns and "c_idx" in fact.columns: print(f"[{GRID_RES_M}m] Using existing r_idx / c_idx from fact table.") else: print(f"[{GRID_RES_M}m] r_idx / c_idx 없음. grid_row / grid_col 기준으로 생성.") row_vals = np.sort(fact["grid_row"].unique()) col_vals = np.sort(fact["grid_col"].unique()) row2idx = {r: i for i, r in enumerate(row_vals)} col2idx = {c: i for i, c in enumerate(col_vals)} fact["r_idx"] = fact["grid_row"].map(row2idx) fact["c_idx"] = fact["grid_col"].map(col2idx) T = int(fact["t_idx"].max()) + 1 R = int(fact["r_idx"].max()) + 1 C = int(fact["c_idx"].max()) + 1 print(f"[{GRID_RES_M}m] T (hours) = {T}, R = {R}, C = {C}") # ---------------- 시간 메타데이터 (time_bias) ---------------- time_bias = time_bias.sort_values("t_idx").drop_duplicates(subset=["t_idx"]) if time_bias["t_idx"].max() + 1 < T: T = int(time_bias["t_idx"].max()) + 1 print(f"[{GRID_RES_M}m] Adjusted T to {T} based on time_bias.") hour_arr = np.zeros(T, dtype=np.int16) dow_arr = np.zeros(T, dtype=np.int16) is_weekend_arr = np.zeros(T, dtype=np.int8) t_idx_meta = time_bias["t_idx"].to_numpy(dtype=int) hours = time_bias["hour"].to_numpy(dtype=int) dows = time_bias["weekday"].to_numpy(dtype=int) weekends = time_bias["is_weekend"].to_numpy(dtype=np.int8) hour_arr[t_idx_meta] = hours dow_arr[t_idx_meta] = dows is_weekend_arr[t_idx_meta] = weekends # ---------------- panel 생성 ---------------- print(f"[{GRID_RES_M}m] Aggregating mean response time per (t_idx, r_idx, c_idx)...") if "소요시각" not in fact.columns: raise ValueError(f"[{GRID_RES_M}m] fact 테이블에 '소요시각' 컬럼이 없습니다.") # (t_idx, r_idx, c_idx)별 평균 소요시각 agg_y = ( fact .groupby(["t_idx", "r_idx", "c_idx"]) .agg( mean_time=("소요시각", "mean"), ) .reset_index() ) # dyn: (grid_id, 단위_m, t_idx)별 cnt_events, mean_dist, GRID_SPEED, link_cnt grid_map = ( fact[["grid_id", "단위_m", "r_idx", "c_idx"]] .drop_duplicates(subset=["grid_id", "단위_m"]) ) dyn = dyn.merge( grid_map, on=["grid_id", "단위_m"], how="left", validate="many_to_one", ) dyn = dyn.dropna(subset=["r_idx", "c_idx"]).copy() dyn["r_idx"] = dyn["r_idx"].astype(int) dyn["c_idx"] = dyn["c_idx"].astype(int) # dyn + mean_time merge panel = dyn.merge( agg_y, on=["t_idx", "r_idx", "c_idx"], how="left", validate="one_to_one", ) # 이름 정리: 모델에서 사용할 이름 panel = panel.rename( columns={ "cnt_events": "cnt", "GRID_SPEED": "traffic", # 동적 traffic 피처 # link_cnt는 지금은 사용 안 함 (원하면 static으로 따로 쓸 수 있음) } ) for col in ["cnt", "mean_dist", "traffic", "mean_time"]: if col not in panel.columns: raise ValueError(f"[{GRID_RES_M}m] panel에 '{col}' 컬럼이 없습니다.") panel[col] = panel[col].fillna(0.0).astype(np.float32) panel = panel[["t_idx", "r_idx", "c_idx", "cnt", "mean_time", "mean_dist", "traffic"]].copy() print(f"[{GRID_RES_M}m] panel rows:", len(panel)) # ---------------- static: (R, C, 3) ---------------- print(f"[{GRID_RES_M}m] Building static grid features...") row_map = fact[["grid_row", "r_idx"]].drop_duplicates() col_map = fact[["grid_col", "c_idx"]].drop_duplicates() row2idx = dict(zip(row_map["grid_row"], row_map["r_idx"])) col2idx = dict(zip(col_map["grid_col"], col_map["c_idx"])) grid_static["r_idx"] = grid_static["grid_row"].map(row2idx) grid_static["c_idx"] = grid_static["grid_col"].map(col2idx) grid_static = grid_static.dropna(subset=["r_idx", "c_idx"]).copy() grid_static["r_idx"] = grid_static["r_idx"].astype(int) grid_static["c_idx"] = grid_static["c_idx"].astype(int) for col in ["road_len_in_grid", "n_links", "n_segments"]: if col not in grid_static.columns: grid_static[col] = 0.0 else: grid_static[col] = grid_static[col].fillna(0.0) features = ["road_len_in_grid", "n_links", "n_segments"] F_stat = len(features) static = np.zeros((R, C, F_stat), dtype=np.float32) r_idx_arr = grid_static["r_idx"].to_numpy(dtype=int) c_idx_arr = grid_static["c_idx"].to_numpy(dtype=int) for i, feat in enumerate(features): vals = grid_static[feat].to_numpy(dtype=np.float32) static[r_idx_arr, c_idx_arr, i] = vals print(f"[{GRID_RES_M}m] static shape:", static.shape) return panel, static, T, R, C, hour_arr, dow_arr, is_weekend_arr