huangfu/photonAI
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

修改脚本结构MOE

Browse Source
This commit is contained in:
huangfu
2026-04-19 18:33:36 +08:00
parent 745868a456
commit 6693f3916c
13 changed files with 1675 additions and 1096 deletions
Download Patch File Download Diff File Expand all files Collapse all files

547
src/preprocess.py
View File

@@ -7,7 +7,7 @@ import logging
import pickle
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from typing import List, Sequence, Tuple
import numpy as np
import pandas as pd
@@ -17,331 +17,55 @@ from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader, TensorDataset
from src.config import AppConfig
from src.data import ALL_COLUMNS, INPUT_COLUMNS, TARGET_COLUMNS, quality_report_before_clean
from src.data import INPUT_COLUMNS, TARGET_COLUMNS, quality_report_before_clean
logger = logging.getLogger(__name__)
@dataclass
class ProcessedDataBundle:
"""训练用张量与 DataLoader，以及划分后的 numpy（含测试集原始物理量用于导出）。"""
"""训练用张量与 DataLoader，以及 notebook 风格的 train/test 划分数据。"""
train_loader: DataLoader
val_loader: DataLoader
test_loader: DataLoader
X_train: np.ndarray
X_val: np.ndarray
X_test: np.ndarray
y_train: np.ndarray
y_val: 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_scaler: StandardScaler
y_scalers: List[StandardScaler]
feature_names: List[str]
target_names: List[str]
def _mask_outliers_iqr(
values: np.ndarray, col_names: List[str], k: float
) -> np.ndarray:
"""返回 True 表示该行在任一选定列上超出训练集 IQR 范围（基于传入的 values 统计）。"""
mask = np.zeros(len(values), dtype=bool)
for j, _ in enumerate(col_names):
col = values[:, j]
q1, q3 = np.percentile(col, [25, 75])
iqr = q3 - q1
lo, hi = q1 - k * iqr, q3 + k * iqr
mask |= (col < lo) | (col > hi)
return mask
def _mask_outliers_zscore(values: np.ndarray, threshold: float) -> np.ndarray:
mask = np.zeros(len(values), dtype=bool)
for j in range(values.shape[1]):
col = values[:, j]
mu, sig = col.mean(), col.std(ddof=0)
if sig < 1e-12:
continue
z = np.abs((col - mu) / sig)
mask |= z > threshold
return mask
def _winsorize_train_apply_all(
train: np.ndarray,
val: np.ndarray,
test: np.ndarray,
ql: float,
qu: float,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""按训练集分位数对 train/val/test 同步裁剪（列方向）。"""
lo = np.quantile(train, ql, axis=0)
hi = np.quantile(train, qu, axis=0)
def clip_arr(a: np.ndarray) -> np.ndarray:
return np.clip(a, lo, hi)
return clip_arr(train), clip_arr(val), clip_arr(test)
def clean_dataframe(
df: pd.DataFrame,
cfg: AppConfig,
report_lines: List[str],
) -> pd.DataFrame:
"""
清洗流程（顺序固定，便于复现与审计）：
1. 可选：完全重复行去重。
2. 可选：以 **最后一列对应字段 V_pi**（txt 第 11 个逗号分隔字段）为门控，仅保留
``v_pi_min <= V_pi <= v_pi_max``（默认 [0, 500]）。
3. 可选：再移除 ``V_pi <= 0``（与区间门控独立，由配置控制）。
"""
"""按 notebook 逻辑清洗：仅保留 V_pi < 阈值。"""
out = df.copy()
n0 = len(out)
if cfg.remove_duplicate_rows:
out = out.drop_duplicates()
report_lines.append(f"去完全重复行: {n0} -> {len(out)}")
if cfg.filter_v_pi_range:
n1 = len(out)
lo, hi = float(cfg.v_pi_min), float(cfg.v_pi_max)
mask = (out["V_pi"] >= lo) & (out["V_pi"] <= hi)
out = out[mask].reset_index(drop=True)
report_lines.append(
f"V_pi 物理区间过滤 [{lo}, {hi}]（txt 第 11 列 / 列名 V_pi）: {n1} -> {len(out)}"
)
if cfg.remove_nonpositive_vpi:
n2 = len(out)
out = out[out["V_pi"] > 0].reset_index(drop=True)
report_lines.append(f"移除 V_pi<=0: {n2} -> {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 区间配置、数据源或是否过度去重。"
)
raise ValueError("清洗后样本数为 0，请检查数据源或 V_pi 阈值设置。")
return out
def _random_split_indices(
n: int,
ratios: List[float],
seed: int,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""返回 train/val/test 的整数索引（先 shuffle 再按比例切分）。"""
rng = np.random.default_rng(seed)
idx = np.arange(n)
rng.shuffle(idx)
tr, va, te = ratios
n_test = int(round(n * te))
n_val = int(round(n * va))
n_train = n - n_val - n_test
if n_train <= 0 or n_val <= 0 or n_test <= 0:
raise ValueError(
f"划分后样本过少: train={n_train}, val={n_val}, test={n_test}，请调整比例或数据量"
)
i_train = idx[:n_train]
i_val = idx[n_train : n_train + n_val]
i_test = idx[n_train + n_val :]
return i_train, i_val, i_test
def _quantile_bin_labels(values: np.ndarray, n_bins: int) -> np.ndarray | None:
"""
基于秩做近似等频分桶，避免重复值导致的 qcut 退化。
返回每个样本所属桶标签；若样本过少则返回 None。
"""
if len(values) < 2:
return None
q = min(int(n_bins), len(values))
if q < 2:
return None
ranks = pd.Series(values).rank(method="first")
labels = pd.qcut(ranks, q=q, labels=False, duplicates="drop")
if labels is None:
return None
arr = np.asarray(labels, dtype=int)
if len(np.unique(arr)) < 2:
return None
return arr
def _grouped_split_indices(
df: pd.DataFrame,
ratios: List[float],
seed: int,
stratify_target: str,
stratify_bins: int,
report_lines: List[str],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
先按输入 8 维分组，再在组级别按目标统计量近似分层切分。
这样可避免“同输入异输出”同时落在 train/val/test，提升评估稳定性。
若分层条件不足，则退化为组级随机切分。
"""
if len(df) < 3:
raise ValueError("样本数过少，无法做 train/val/test 切分。")
group_ids = df.groupby(INPUT_COLUMNS, sort=False, dropna=False).ngroup().to_numpy()
n_groups = int(group_ids.max()) + 1
group_df = df.copy()
group_df["_group_id"] = group_ids
group_stat = (
group_df.groupby("_group_id", sort=True)
.agg(group_size=("V_pi", "size"), strat_value=(stratify_target, "median"))
.reset_index()
)
group_id_arr = group_stat["_group_id"].to_numpy(dtype=int)
labels = _quantile_bin_labels(
group_stat["strat_value"].to_numpy(dtype=np.float64),
stratify_bins,
)
tr, va, te = ratios
holdout_ratio = va + te
val_ratio_in_holdout = va / holdout_ratio
def _split_groups(use_stratify: bool) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
strat = labels if use_stratify and labels is not None else None
train_groups, holdout_groups = train_test_split(
group_id_arr,
train_size=tr,
test_size=holdout_ratio,
random_state=seed,
shuffle=True,
stratify=strat,
)
holdout_strat = None
if strat is not None:
label_map = dict(zip(group_id_arr.tolist(), labels.tolist()))
holdout_labels = np.asarray(
[label_map[int(g)] for g in holdout_groups],
dtype=int,
)
if len(np.unique(holdout_labels)) >= 2:
holdout_strat = holdout_labels
val_groups, test_groups = train_test_split(
holdout_groups,
train_size=val_ratio_in_holdout,
test_size=1.0 - val_ratio_in_holdout,
random_state=seed + 1,
shuffle=True,
stratify=holdout_strat,
)
return (
np.asarray(train_groups, dtype=int),
np.asarray(val_groups, dtype=int),
np.asarray(test_groups, dtype=int),
)
split_note = (
f"按输入分组切分，共 {n_groups} 个唯一输入组；"
f"组级按 {stratify_target} 中位数分 {min(stratify_bins, n_groups)} 桶近似分层。"
)
try:
train_groups, val_groups, test_groups = _split_groups(use_stratify=True)
report_lines.append(split_note)
except ValueError as e:
train_groups, val_groups, test_groups = _split_groups(use_stratify=False)
report_lines.append(f"{split_note} 但分层条件不足，退化为组级随机切分：{e}")
i_train = np.flatnonzero(np.isin(group_ids, train_groups))
i_val = np.flatnonzero(np.isin(group_ids, val_groups))
i_test = np.flatnonzero(np.isin(group_ids, test_groups))
return i_train, i_val, i_test
def build_split_indices(
df: pd.DataFrame,
cfg: AppConfig,
report_lines: List[str],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""根据配置生成 train/val/test 行索引。"""
if cfg.split_mode == "random":
report_lines.append("切分策略：随机打乱后按比例切分。")
return _random_split_indices(len(df), cfg.split_ratios, cfg.random_seed)
return _grouped_split_indices(
df=df,
ratios=cfg.split_ratios,
seed=cfg.random_seed,
stratify_target=cfg.split_stratify_target,
stratify_bins=cfg.split_stratify_bins,
report_lines=report_lines,
)
def apply_train_only_outliers(
X_train: np.ndarray,
y_train: np.ndarray,
X_val: np.ndarray,
y_val: np.ndarray,
X_test: np.ndarray,
y_test: np.ndarray,
cfg: AppConfig,
report_lines: List[str],
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""
仅在训练集上估计阈值：
- iqr/zscore: 从训练集删除离群行（val/test 不动）
- quantile_clip: 对 train/val/test 同步 winsorize（阈值来自 train）
"""
strat = cfg.outlier_strategy
if strat == "none":
report_lines.append("outlier_strategy=none：不对数值做裁剪/删除（除配置项外）。")
return X_train, y_train, X_val, y_val, X_test, y_test
cols = cfg.outlier_apply_to
if cols == "all":
train_mat = np.hstack([X_train, y_train])
val_mat = np.hstack([X_val, y_val])
test_mat = np.hstack([X_test, y_test])
names = INPUT_COLUMNS + TARGET_COLUMNS
else:
train_mat = y_train.copy()
val_mat = y_val.copy()
test_mat = y_test.copy()
names = TARGET_COLUMNS
if strat == "quantile_clip":
ql = cfg.outlier_config.quantile_lower
qu = cfg.outlier_config.quantile_upper
tr2, va2, te2 = _winsorize_train_apply_all(train_mat, val_mat, test_mat, ql, qu)
report_lines.append(
f"quantile_clip: 按训练集分位数 [{ql}, {qu}] 对 {cols} 列 winsorize。"
)
if cols == "all":
d = len(INPUT_COLUMNS)
X_train, y_train = tr2[:, :d], tr2[:, d:]
X_val, y_val = va2[:, :d], va2[:, d:]
X_test, y_test = te2[:, :d], te2[:, d:]
else:
y_train, y_val, y_test = tr2, va2, te2
return X_train, y_train, X_val, y_val, X_test, y_test
if strat == "iqr":
mask = _mask_outliers_iqr(train_mat, names, cfg.outlier_config.iqr_k)
elif strat == "zscore":
mask = _mask_outliers_zscore(train_mat, cfg.outlier_config.zscore_threshold)
else:
raise ValueError(f"未知 outlier_strategy: {strat}")
removed = int(mask.sum())
kept = ~mask
X_train, y_train = X_train[kept], y_train[kept]
report_lines.append(
f"{strat}: 在训练子集上检测 {cols} 离群，删除训练行 {removed}，保留 {len(X_train)}。"
)
return X_train, y_train, X_val, y_val, X_test, y_test
def build_dataloaders(
X_train: np.ndarray,
y_train: np.ndarray,
X_val: np.ndarray,
y_val: np.ndarray,
X_test: np.ndarray,
y_test: np.ndarray,
batch_size: int,
num_workers: int,
) -> Tuple[DataLoader, DataLoader, DataLoader]:
) -> Tuple[DataLoader, DataLoader]:
def to_loader(X: np.ndarray, y: np.ndarray, shuffle: bool) -> DataLoader:
ds = TensorDataset(
torch.from_numpy(X).float(),
@@ -357,127 +81,77 @@ def build_dataloaders(
return (
to_loader(X_train, y_train, shuffle=True),
to_loader(X_val, y_val, shuffle=False),
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_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_scaler.pkl").open("wb") as f:
pickle.dump(y_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, StandardScaler]:
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_scaler.pkl").open("rb") as f:
y_scaler = pickle.load(f)
return X_scaler, y_scaler
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_val: np.ndarray,
i_test: np.ndarray,
) -> None:
"""保存对清洗后矩阵行的划分索引，便于 eval 阶段完全复现。"""
payload = {
"train": i_train.astype(int).tolist(),
"val": i_val.astype(int).tolist(),
"test": i_test.astype(int).tolist(),
}
(run_dir / "split_indices.json").write_text(
json.dumps(payload, indent=2), encoding="utf-8"
json.dumps(payload, indent=2),
encoding="utf-8",
)
def load_split_indices(run_dir: Path) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
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}。请使用本仓库训练产生的 run 目录，或先完成一次训练。"
)
raise FileNotFoundError(f"未找到 {path}")
data = json.loads(path.read_text(encoding="utf-8"))
return (
np.asarray(data["train"], dtype=int),
np.asarray(data["val"], dtype=int),
np.asarray(data["test"], dtype=int),
)
def rebuild_bundle_for_eval(
df: pd.DataFrame,
cfg: AppConfig,
run_dir: Path,
) -> ProcessedDataBundle:
"""
与训练阶段相同的清洗、划分与离群处理，但使用已保存的 StandardScaler 仅做 transform。
用于独立 eval / infer 流程，避免重新拟合 scaler 造成分布偏移。
"""
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_tr, i_va, i_te = load_split_indices(run_dir)
for name, idx in ("train", i_tr), ("val", i_va), ("test", i_te):
if len(idx) == 0 or int(idx.max()) >= len(X) or int(idx.min()) < 0:
raise ValueError(
f"split_indices.json 与当前数据不兼容（{name} 索引越界或为空）。"
f"请确认 data_path 指向与训练相同的清洗后样本空间。"
)
X_train, y_train = X[i_tr], y[i_tr]
X_val, y_val = X[i_va], y[i_va]
X_test, y_test = X[i_te], y[i_te]
X_train, y_train, X_val, y_val, X_test, y_test = apply_train_only_outliers(
X_train, y_train, X_val, y_val, X_test, y_test, cfg, report_lines
)
X_scaler, y_scaler = load_scalers(run_dir)
X_train_s = X_scaler.transform(X_train)
y_train_s = y_scaler.transform(y_train)
X_val_s = X_scaler.transform(X_val)
y_val_s = y_scaler.transform(y_val)
X_test_s = X_scaler.transform(X_test)
y_test_s = y_scaler.transform(y_test)
train_loader, val_loader, test_loader = build_dataloaders(
X_train_s,
y_train_s,
X_val_s,
y_val_s,
X_test_s,
y_test_s,
cfg.training.batch_size,
cfg.training.num_workers,
)
return ProcessedDataBundle(
train_loader=train_loader,
val_loader=val_loader,
test_loader=test_loader,
X_train=X_train_s,
X_val=X_val_s,
X_test=X_test_s,
y_train=y_train_s,
y_val=y_val_s,
y_test=y_test_s,
X_test_raw=X_test,
y_test_raw=y_test,
X_scaler=X_scaler,
y_scaler=y_scaler,
feature_names=list(INPUT_COLUMNS),
target_names=list(TARGET_COLUMNS),
)
def write_data_report_md(
path: Path,
raw_quality: dict,
@@ -515,12 +189,9 @@ def prepare_training_data(
cfg: AppConfig,
run_dir: Path,
) -> ProcessedDataBundle:
"""
完整预处理流水线：质量报告 -> 清洗 -> 划分 -> 训练集离群处理 -> 标准化 -> DataLoader。
将 data_report.md 与 cleaning 元数据写入 run_dir。
"""
"""按 notebook 一致逻辑准备 train/test、scaler 与 DataLoader。"""
report_lines: List[str] = []
raw_q = quality_report_before_clean(df, cfg.v_pi_min, cfg.v_pi_max)
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)
@@ -528,39 +199,35 @@ def prepare_training_data(
X = cleaned[INPUT_COLUMNS].to_numpy(dtype=np.float64)
y = cleaned[TARGET_COLUMNS].to_numpy(dtype=np.float64)
all_idx = np.arange(len(X))
i_tr, i_va, i_te = build_split_indices(cleaned, cfg, report_lines)
save_split_indices(run_dir, i_tr, i_va, i_te)
X_train, y_train = X[i_tr], y[i_tr]
X_val, y_val = X[i_va], y[i_va]
X_test, y_test = X[i_te], y[i_te]
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/val/test = {cfg.split_ratios}，样本数 "
f"{len(X_train)}/{len(X_val)}/{len(X_test)}"
)
X_train, y_train, X_val, y_val, X_test, y_test = apply_train_only_outliers(
X_train, y_train, X_val, y_val, X_test, y_test, cfg, report_lines
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()
y_scaler = StandardScaler()
X_train_s = X_scaler.fit_transform(X_train)
y_train_s = y_scaler.fit_transform(y_train)
X_val_s = X_scaler.transform(X_val)
y_val_s = y_scaler.transform(y_val)
X_test_s = X_scaler.transform(X_test)
y_test_s = y_scaler.transform(y_test)
save_scalers(X_scaler, y_scaler, run_dir)
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,
"split_ratios": cfg.split_ratios,
"n_train": int(len(X_train_s)),
"n_val": int(len(X_val_s)),
"n_test": int(len(X_test_s)),
"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),
@@ -576,31 +243,77 @@ def prepare_training_data(
stats_after.to_csv(run_dir / "data_stats.csv", encoding="utf-8")
logger.info("预处理完成：%s", run_dir / "data_report.md")
train_loader, val_loader, test_loader = build_dataloaders(
X_train_s,
y_train_s,
X_val_s,
y_val_s,
X_test_s,
y_test_s,
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,
val_loader=val_loader,
test_loader=test_loader,
X_train=X_train_s,
X_val=X_val_s,
X_test=X_test_s,
y_train=y_train_s,
y_val=y_val_s,
y_test=y_test_s,
X_test_raw=X_test,
y_test_raw=y_test,
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_scaler=y_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),
)