320 lines
9.5 KiB
Python
320 lines
9.5 KiB
Python
"""清洗、划分、标准化与 DataLoader 构建。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import pickle
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
from typing import List, Sequence, Tuple
|
||
|
||
import numpy as np
|
||
import pandas as pd
|
||
import torch
|
||
from sklearn.model_selection import train_test_split
|
||
from sklearn.preprocessing import StandardScaler
|
||
from torch.utils.data import DataLoader, TensorDataset
|
||
|
||
from src.config import AppConfig
|
||
from src.data import INPUT_COLUMNS, TARGET_COLUMNS, quality_report_before_clean
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class ProcessedDataBundle:
|
||
"""训练用张量与 DataLoader，以及 notebook 风格的 train/test 划分数据。"""
|
||
|
||
train_loader: DataLoader
|
||
test_loader: DataLoader
|
||
X_train: np.ndarray
|
||
X_test: np.ndarray
|
||
y_train: np.ndarray
|
||
y_test: np.ndarray
|
||
X_train_raw: np.ndarray
|
||
X_test_raw: np.ndarray
|
||
y_train_raw: np.ndarray
|
||
y_test_raw: np.ndarray
|
||
X_scaler: StandardScaler
|
||
y_scalers: List[StandardScaler]
|
||
feature_names: List[str]
|
||
target_names: List[str]
|
||
|
||
|
||
def clean_dataframe(
|
||
df: pd.DataFrame,
|
||
cfg: AppConfig,
|
||
report_lines: List[str],
|
||
) -> pd.DataFrame:
|
||
"""按 notebook 逻辑清洗：仅保留 V_pi < 阈值。"""
|
||
out = df.copy()
|
||
n0 = len(out)
|
||
vmax = float(cfg.data.filter_v_pi_max)
|
||
out = out[out["V_pi"] < vmax].reset_index(drop=True)
|
||
report_lines.append(f"V_pi 阈值过滤 (< {vmax})：{n0} -> {len(out)}")
|
||
if len(out) == 0:
|
||
raise ValueError("清洗后样本数为 0，请检查数据源或 V_pi 阈值设置。")
|
||
return out
|
||
|
||
|
||
def build_dataloaders(
|
||
X_train: np.ndarray,
|
||
y_train: np.ndarray,
|
||
X_test: np.ndarray,
|
||
y_test: np.ndarray,
|
||
batch_size: int,
|
||
num_workers: int,
|
||
) -> Tuple[DataLoader, DataLoader]:
|
||
def to_loader(X: np.ndarray, y: np.ndarray, shuffle: bool) -> DataLoader:
|
||
ds = TensorDataset(
|
||
torch.from_numpy(X).float(),
|
||
torch.from_numpy(y).float(),
|
||
)
|
||
return DataLoader(
|
||
ds,
|
||
batch_size=batch_size,
|
||
shuffle=shuffle,
|
||
num_workers=num_workers,
|
||
pin_memory=False,
|
||
)
|
||
|
||
return (
|
||
to_loader(X_train, y_train, shuffle=True),
|
||
to_loader(X_test, y_test, shuffle=False),
|
||
)
|
||
|
||
|
||
def _fit_target_scalers(y_train: np.ndarray) -> Tuple[List[StandardScaler], np.ndarray]:
|
||
scalers: List[StandardScaler] = []
|
||
scaled_cols = []
|
||
for i in range(y_train.shape[1]):
|
||
scaler = StandardScaler()
|
||
scaled_cols.append(scaler.fit_transform(y_train[:, i : i + 1]))
|
||
scalers.append(scaler)
|
||
return scalers, np.hstack(scaled_cols)
|
||
|
||
|
||
def transform_targets(y: np.ndarray, y_scalers: Sequence[StandardScaler]) -> np.ndarray:
|
||
cols = [scaler.transform(y[:, i : i + 1]) for i, scaler in enumerate(y_scalers)]
|
||
return np.hstack(cols)
|
||
|
||
|
||
def inverse_transform_targets(
|
||
y_scaled: np.ndarray, y_scalers: Sequence[StandardScaler]
|
||
) -> np.ndarray:
|
||
cols = [scaler.inverse_transform(y_scaled[:, i : i + 1]) for i, scaler in enumerate(y_scalers)]
|
||
return np.hstack(cols)
|
||
|
||
|
||
def save_scalers(
|
||
X_scaler: StandardScaler,
|
||
y_scalers: Sequence[StandardScaler],
|
||
run_dir: Path,
|
||
) -> None:
|
||
with (run_dir / "x_scaler.pkl").open("wb") as f:
|
||
pickle.dump(X_scaler, f)
|
||
with (run_dir / "y_scalers.pkl").open("wb") as f:
|
||
pickle.dump(list(y_scalers), f)
|
||
|
||
|
||
def load_scalers(run_dir: Path) -> Tuple[StandardScaler, List[StandardScaler]]:
|
||
with (run_dir / "x_scaler.pkl").open("rb") as f:
|
||
X_scaler = pickle.load(f)
|
||
with (run_dir / "y_scalers.pkl").open("rb") as f:
|
||
y_scalers = pickle.load(f)
|
||
return X_scaler, list(y_scalers)
|
||
|
||
|
||
def save_split_indices(
|
||
run_dir: Path,
|
||
i_train: np.ndarray,
|
||
i_test: np.ndarray,
|
||
) -> None:
|
||
payload = {
|
||
"train": i_train.astype(int).tolist(),
|
||
"test": i_test.astype(int).tolist(),
|
||
}
|
||
(run_dir / "split_indices.json").write_text(
|
||
json.dumps(payload, indent=2),
|
||
encoding="utf-8",
|
||
)
|
||
|
||
|
||
def load_split_indices(run_dir: Path) -> Tuple[np.ndarray, np.ndarray]:
|
||
path = run_dir / "split_indices.json"
|
||
if not path.is_file():
|
||
raise FileNotFoundError(f"未找到 {path}")
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
return (
|
||
np.asarray(data["train"], dtype=int),
|
||
np.asarray(data["test"], dtype=int),
|
||
)
|
||
|
||
|
||
def write_data_report_md(
|
||
path: Path,
|
||
raw_quality: dict,
|
||
report_lines: List[str],
|
||
basic_stats_before: pd.DataFrame,
|
||
basic_stats_after: pd.DataFrame,
|
||
) -> None:
|
||
lines = [
|
||
"# 数据与清洗报告",
|
||
"",
|
||
"## 清洗前质量摘要（JSON）",
|
||
"```json",
|
||
json.dumps(raw_quality, indent=2, ensure_ascii=False, default=str),
|
||
"```",
|
||
"",
|
||
"## 清洗步骤",
|
||
"\n".join(f"- {x}" for x in report_lines),
|
||
"",
|
||
"## 清洗前 describe（CSV 文本块）",
|
||
"```text",
|
||
basic_stats_before.to_csv(),
|
||
"```",
|
||
"",
|
||
"## 清洗后 describe（CSV 文本块）",
|
||
"```text",
|
||
basic_stats_after.to_csv(),
|
||
"```",
|
||
"",
|
||
]
|
||
path.write_text("\n".join(lines), encoding="utf-8")
|
||
|
||
|
||
def prepare_training_data(
|
||
df: pd.DataFrame,
|
||
cfg: AppConfig,
|
||
run_dir: Path,
|
||
) -> ProcessedDataBundle:
|
||
"""按 notebook 一致逻辑准备 train/test、scaler 与 DataLoader。"""
|
||
report_lines: List[str] = []
|
||
raw_q = quality_report_before_clean(df, 0.0, cfg.data.filter_v_pi_max)
|
||
stats_before = df.describe().T
|
||
|
||
cleaned = clean_dataframe(df, cfg, report_lines)
|
||
stats_after = cleaned.describe().T
|
||
|
||
X = cleaned[INPUT_COLUMNS].to_numpy(dtype=np.float64)
|
||
y = cleaned[TARGET_COLUMNS].to_numpy(dtype=np.float64)
|
||
all_idx = np.arange(len(X))
|
||
|
||
X_train_raw, X_test_raw, y_train_raw, y_test_raw, i_train, i_test = train_test_split(
|
||
X,
|
||
y,
|
||
all_idx,
|
||
test_size=cfg.data.test_size,
|
||
random_state=cfg.data.random_state,
|
||
)
|
||
report_lines.append(
|
||
f"train_test_split(test_size={cfg.data.test_size}, random_state={cfg.data.random_state}) "
|
||
f"-> {len(X_train_raw)}/{len(X_test_raw)}"
|
||
)
|
||
save_split_indices(run_dir, i_train, i_test)
|
||
|
||
X_scaler = StandardScaler()
|
||
X_train = X_scaler.fit_transform(X_train_raw)
|
||
X_test = X_scaler.transform(X_test_raw)
|
||
y_scalers, y_train = _fit_target_scalers(y_train_raw)
|
||
y_test = transform_targets(y_test_raw, y_scalers)
|
||
save_scalers(X_scaler, y_scalers, run_dir)
|
||
|
||
meta = {
|
||
"raw_quality": raw_q,
|
||
"cleaning_steps": report_lines,
|
||
"test_size": cfg.data.test_size,
|
||
"random_state": cfg.data.random_state,
|
||
"n_train": int(len(X_train)),
|
||
"n_test": int(len(X_test)),
|
||
}
|
||
(run_dir / "cleaning_meta.json").write_text(
|
||
json.dumps(meta, indent=2, ensure_ascii=False, default=str),
|
||
encoding="utf-8",
|
||
)
|
||
write_data_report_md(
|
||
run_dir / "data_report.md",
|
||
raw_q,
|
||
report_lines,
|
||
stats_before,
|
||
stats_after,
|
||
)
|
||
stats_after.to_csv(run_dir / "data_stats.csv", encoding="utf-8")
|
||
logger.info("预处理完成：%s", run_dir / "data_report.md")
|
||
|
||
train_loader, test_loader = build_dataloaders(
|
||
X_train,
|
||
y_train,
|
||
X_test,
|
||
y_test,
|
||
cfg.training.batch_size,
|
||
cfg.training.num_workers,
|
||
)
|
||
|
||
return ProcessedDataBundle(
|
||
train_loader=train_loader,
|
||
test_loader=test_loader,
|
||
X_train=X_train,
|
||
X_test=X_test,
|
||
y_train=y_train,
|
||
y_test=y_test,
|
||
X_train_raw=X_train_raw,
|
||
X_test_raw=X_test_raw,
|
||
y_train_raw=y_train_raw,
|
||
y_test_raw=y_test_raw,
|
||
X_scaler=X_scaler,
|
||
y_scalers=y_scalers,
|
||
feature_names=list(INPUT_COLUMNS),
|
||
target_names=list(TARGET_COLUMNS),
|
||
)
|
||
|
||
|
||
def rebuild_bundle_for_eval(
|
||
df: pd.DataFrame,
|
||
cfg: AppConfig,
|
||
run_dir: Path,
|
||
) -> ProcessedDataBundle:
|
||
"""使用训练时保存的切分索引与 scaler 重新构造 train/test 数据。"""
|
||
report_lines: List[str] = []
|
||
cleaned = clean_dataframe(df, cfg, report_lines)
|
||
X = cleaned[INPUT_COLUMNS].to_numpy(dtype=np.float64)
|
||
y = cleaned[TARGET_COLUMNS].to_numpy(dtype=np.float64)
|
||
i_train, i_test = load_split_indices(run_dir)
|
||
for name, idx in (("train", i_train), ("test", i_test)):
|
||
if len(idx) == 0 or int(idx.max()) >= len(X) or int(idx.min()) < 0:
|
||
raise ValueError(f"split_indices.json 与当前数据不兼容（{name} 索引越界或为空）")
|
||
|
||
X_train_raw, X_test_raw = X[i_train], X[i_test]
|
||
y_train_raw, y_test_raw = y[i_train], y[i_test]
|
||
X_scaler, y_scalers = load_scalers(run_dir)
|
||
X_train = X_scaler.transform(X_train_raw)
|
||
X_test = X_scaler.transform(X_test_raw)
|
||
y_train = transform_targets(y_train_raw, y_scalers)
|
||
y_test = transform_targets(y_test_raw, y_scalers)
|
||
train_loader, test_loader = build_dataloaders(
|
||
X_train,
|
||
y_train,
|
||
X_test,
|
||
y_test,
|
||
cfg.training.batch_size,
|
||
cfg.training.num_workers,
|
||
)
|
||
|
||
return ProcessedDataBundle(
|
||
train_loader=train_loader,
|
||
test_loader=test_loader,
|
||
X_train=X_train,
|
||
X_test=X_test,
|
||
y_train=y_train,
|
||
y_test=y_test,
|
||
X_train_raw=X_train_raw,
|
||
X_test_raw=X_test_raw,
|
||
y_train_raw=y_train_raw,
|
||
y_test_raw=y_test_raw,
|
||
X_scaler=X_scaler,
|
||
y_scalers=y_scalers,
|
||
feature_names=list(INPUT_COLUMNS),
|
||
target_names=list(TARGET_COLUMNS),
|
||
)
|