# -*- coding: utf-8 -*- """原始 HDF5 样本预处理。 数据生成阶段写出的 HDF5 文件仍处在“原始样本”状态：参数尚未做特征变换，输入和 监督目标尚未标准化，也没有训练/验证/测试划分。本模块完成这些准备工作，并把 结果保存为 joblib 文件供训练脚本直接读取。 预处理只在训练集上拟合 scaler，随后复用到验证集和测试集，避免数据泄漏。输出 payload 中保留 curve_layout、参数变换配置、流量制度元数据和附加字段，使后续 训练、评估和误差分析可以追溯每个样本的来源。 """ from __future__ import annotations import copy from pathlib import Path import h5py import joblib import numpy as np from sklearn.model_selection import GroupShuffleSplit, train_test_split from sklearn.preprocessing import StandardScaler from src.data.param_features import build_param_feature_transform, transform_param_features class SelectiveStandardScaler: """只标准化指定连续列，同时保持 one-hot 等类别列原值不变。""" def __init__(self, scaled_indices: np.ndarray | list[int]): self.scaled_indices = np.asarray(scaled_indices, dtype=np.int64) def fit(self, x: np.ndarray) -> "SelectiveStandardScaler": """在训练集的连续特征列上拟合均值和标准差。""" values = _ensure_2d("params", x).astype(np.float64) indices = np.unique(self.scaled_indices) if np.any(indices < 0) or np.any(indices >= values.shape[1]): raise ValueError(f"scaled parameter indices out of range: {indices.tolist()}") self.scaled_indices_ = indices self.n_features_in_ = int(values.shape[1]) self.n_samples_seen_ = int(values.shape[0]) self.mean_ = np.zeros(self.n_features_in_, dtype=np.float64) self.var_ = np.ones(self.n_features_in_, dtype=np.float64) self.scale_ = np.ones(self.n_features_in_, dtype=np.float64) if indices.size: inner = StandardScaler().fit(values[:, indices]) self.mean_[indices] = inner.mean_ self.var_[indices] = inner.var_ self.scale_[indices] = inner.scale_ return self def _validate_transform_input(self, x: np.ndarray) -> np.ndarray: """检查待变换数组维度是否与拟合时一致。""" if not hasattr(self, "n_features_in_"): raise RuntimeError("SelectiveStandardScaler is not fitted") values = _ensure_2d("params", x).astype(np.float64) if values.shape[1] != self.n_features_in_: raise ValueError( f"parameter feature dimension mismatch: {values.shape[1]} != {self.n_features_in_}" ) return values def transform(self, x: np.ndarray) -> np.ndarray: """标准化连续列；未选择的类别列保持0/1。""" values = self._validate_transform_input(x) return (values - self.mean_) / self.scale_ def inverse_transform(self, x: np.ndarray) -> np.ndarray: """恢复标准化前的连续特征，同时保持类别列不变。""" values = self._validate_transform_input(x) return values * self.scale_ + self.mean_ def fit_transform(self, x: np.ndarray) -> np.ndarray: """拟合后立即变换训练集。""" return self.fit(x).transform(x) class MaskedStandardScaler: """每个曲线列只使用实际物理时间覆盖范围内的有效值进行标准化。""" def fit(self, x: np.ndarray, valid_mask: np.ndarray) -> "MaskedStandardScaler": values = _ensure_2d("curve", x).astype(np.float64) mask = np.asarray(valid_mask, dtype=bool) if mask.shape != values.shape: raise ValueError(f"curve mask shape mismatch: {mask.shape} != {values.shape}") if not np.isfinite(values[mask]).all(): raise ValueError("valid curve values contain non-finite values") counts = mask.sum(axis=0).astype(np.int64) sums = np.where(mask, values, 0.0).sum(axis=0) means = np.divide(sums, counts, out=np.zeros_like(sums), where=counts > 0) centered = np.where(mask, values - means, 0.0) variances = np.divide( (centered * centered).sum(axis=0), counts, out=np.ones_like(sums), where=counts > 0, ) scales = np.sqrt(np.maximum(variances, 0.0)) scales[(counts == 0) | (scales < 1.0e-12)] = 1.0 self.n_features_in_ = int(values.shape[1]) self.n_samples_seen_ = counts self.mean_ = means self.var_ = variances self.scale_ = scales return self def _validate(self, x: np.ndarray) -> np.ndarray: if not hasattr(self, "n_features_in_"): raise RuntimeError("MaskedStandardScaler is not fitted") values = _ensure_2d("curve", x).astype(np.float64) if values.shape[1] != self.n_features_in_: raise ValueError( f"curve feature dimension mismatch: {values.shape[1]} != {self.n_features_in_}" ) return values def transform(self, x: np.ndarray) -> np.ndarray: values = self._validate(x) return (values - self.mean_) / self.scale_ def inverse_transform(self, x: np.ndarray) -> np.ndarray: values = self._validate(x) return values * self.scale_ + self.mean_ def _value_counts(values: np.ndarray | None, indices: np.ndarray) -> dict[str, int]: """按字符串键记录某个划分中的类别样本数。""" if values is None: return {} unique, counts = np.unique(np.asarray(values)[indices], return_counts=True) result = {} for value, count in zip(unique.tolist(), counts.tolist()): key = str(int(value)) if float(value).is_integer() else str(value) result[key] = int(count) return result def _split_sample_indices( n_samples: int, test_size: float, val_size: float, random_seed: int, group_id: np.ndarray | None, solver_type: np.ndarray | None, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, str]: """优先按group整体划分；没有可靠group时按solverType分层随机划分。""" if test_size <= 0.0 or val_size <= 0.0 or test_size + val_size >= 1.0: raise ValueError("test_size and val_size must be positive and sum to less than 1") idx = np.arange(n_samples) val_ratio_in_train_val = val_size / (1.0 - test_size) groups = None if group_id is None else np.asarray(group_id).reshape(-1) use_groups = bool( groups is not None and len(groups) == n_samples and np.all(np.isfinite(groups)) and np.all(groups >= 0) and np.unique(groups).size >= 3 ) if use_groups: first = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=random_seed) train_val_pos, test_pos = next(first.split(idx, groups=groups)) idx_train_val = idx[train_val_pos] idx_test = idx[test_pos] second = GroupShuffleSplit( n_splits=1, test_size=val_ratio_in_train_val, random_state=random_seed + 1, ) train_pos, val_pos = next( second.split(idx_train_val, groups=groups[idx_train_val]) ) idx_train = idx_train_val[train_pos] idx_val = idx_train_val[val_pos] return idx_train, idx_val, idx_test, "group_id" stratify = None if solver_type is None else np.asarray(solver_type).reshape(-1) idx_train_val, idx_test = train_test_split( idx, test_size=test_size, random_state=random_seed, shuffle=True, stratify=stratify, ) train_val_stratify = None if stratify is None else stratify[idx_train_val] idx_train, idx_val = train_test_split( idx_train_val, test_size=val_ratio_in_train_val, random_state=random_seed, shuffle=True, stratify=train_val_stratify, ) return idx_train, idx_val, idx_test, "solverType" if stratify is not None else "random" def _ensure_2d(name: str, arr: np.ndarray) -> np.ndarray: """保证数组至少为二维，避免单样本情况下维度被压扁。""" arr = np.asarray(arr, dtype=np.float32) if arr.ndim != 2: raise ValueError(f"{name} must be a 2D array, got shape={arr.shape}") return arr def _decode_attr_list(raw_names) -> list[str] | None: """解码 HDF5 属性中保存的 JSON 字符串或字节串列表。""" if raw_names is None: return None return [ item.decode("utf-8") if isinstance(item, (bytes, np.bytes_)) else str(item) for item in raw_names ] def _read_optional_string_dataset(f: h5py.File, name: str) -> np.ndarray | None: """读取可选字符串数据集；不存在时返回 None。""" if name not in f: return None return np.asarray(f[name][:]).astype(str) def _read_optional_numeric_dataset(f: h5py.File, name: str) -> np.ndarray | None: """读取可选数值数据集；不存在时返回 None。""" if name not in f: return None return np.asarray(f[name][:]) def _validate_optional_length(name: str, arr: np.ndarray | None, expected_n: int) -> None: """检查可选数据集长度是否与主样本数量一致。""" if arr is not None and len(arr) != expected_n: raise ValueError(f"{name} row count does not match main arrays: {len(arr)} != {expected_n}") def _validate_curve_time( curve_time: np.ndarray, expected_n: int, expected_n_time_points: int, ) -> np.ndarray: """校验每条曲线对应的物理时间坐标。""" values = np.asarray(curve_time, dtype=np.float64) expected_shape = (expected_n, expected_n_time_points) if values.ndim != 2 or values.shape != expected_shape: raise ValueError( f"curve_time must have shape {expected_shape}, got shape={values.shape}" ) if not np.all(np.isfinite(values)): raise ValueError("curve_time contains non-finite values") if np.any(values <= 0.0): raise ValueError("curve_time values must be positive") if np.any(np.diff(values, axis=1) <= 0.0): raise ValueError("each curve_time row must be strictly increasing") return values def _validate_curve_valid_mask( curve_valid_mask: np.ndarray, expected_n: int, expected_n_time_points: int, ) -> np.ndarray: """校验逐时间点记录物理覆盖范围的有效掩码。""" values = np.asarray(curve_valid_mask) expected_shape = (expected_n, expected_n_time_points) if values.ndim != 2 or values.shape != expected_shape: raise ValueError( f"curve_valid_mask must have shape {expected_shape}, got shape={values.shape}" ) if not np.all(np.isfinite(values)) or not np.all((values == 0) | (values == 1)): raise ValueError("curve_valid_mask must contain only 0/1 values") values = values.astype(bool) if np.any(values.sum(axis=1) < 2): raise ValueError("each curve_valid_mask row must contain at least two valid points") return values def _curve_keep_indices(meta: dict, curve_dim: int) -> np.ndarray: """返回旧 processed 曲线中 pressure/derivative 两段的列索引。""" layout = meta.get("curve_layout") or {} parts = {str(part["name"]): part for part in layout.get("parts", [])} if "log_pressure" in parts and "log_derivative" in parts: indices = [] for name in ("log_pressure", "log_derivative"): part = parts[name] indices.extend(range(int(part["start"]), int(part["end"]))) return np.asarray(indices, dtype=np.int64) if curve_dim % 3 != 0: raise ValueError( "processed 数据没有 curve_layout，且曲线维度不能按旧版三通道布局识别" ) n_time_points = curve_dim // 3 return np.arange(2 * n_time_points, dtype=np.int64) def _slice_standard_scaler(scaler: StandardScaler, keep_indices: np.ndarray) -> StandardScaler: """裁剪按列独立拟合的 StandardScaler，同时保留其拟合状态。""" result = copy.deepcopy(scaler) for name in ("mean_", "scale_", "var_"): value = getattr(result, name, None) if value is not None: setattr(result, name, np.asarray(value)[keep_indices].copy()) result.n_features_in_ = int(keep_indices.size) return result def migrate_processed_dataset_without_slope(input_path: Path, output_path: Path) -> None: """把旧三通道 processed pkl 精确迁移为 pressure/derivative 双通道数据。""" if not input_path.exists(): raise FileNotFoundError(f"Input file not found: {input_path}") source = joblib.load(input_path) required = ["Y_curve_train", "Y_curve_val", "Y_curve_test", "scaler_curve", "meta"] missing = [name for name in required if name not in source] if missing: raise KeyError(f"processed 数据缺少字段: {missing}") curve_dim = int(source["Y_curve_train"].shape[1]) keep_indices = _curve_keep_indices(source.get("meta", {}), curve_dim) if keep_indices.size == curve_dim: raise ValueError("输入 processed 数据已经不包含 slope，无需迁移") if keep_indices.size % 2 != 0: raise ValueError(f"保留后的曲线维度必须能被 2 整除，当前为 {keep_indices.size}") migration_n_time_points = int(keep_indices.size // 2) migration_meta = dict(source.get("meta", {})) if migration_meta.get("curve_time_mode") != "fixed": raise ValueError( "processed migration requires meta.curve_time_mode='fixed'; " "regenerate the H5 dataset on a fixed time grid and preprocess it again" ) prediction_curve_time = np.asarray( migration_meta.get("prediction_curve_time"), dtype=np.float64, ) if prediction_curve_time.shape != (migration_n_time_points,): raise ValueError( "processed migration requires meta.prediction_curve_time with shape " f"({migration_n_time_points},); regenerate the H5 dataset and preprocess it again" ) if ( not np.all(np.isfinite(prediction_curve_time)) or np.any(prediction_curve_time <= 0.0) or np.any(np.diff(prediction_curve_time) <= 0.0) ): raise ValueError( "processed migration requires a finite, positive, strictly increasing " "meta.prediction_curve_time" ) for split in ("train", "val", "test"): mask_key = f"curve_valid_mask_{split}" if mask_key not in source: raise ValueError( f"processed migration requires {mask_key}; regenerate the H5 dataset " "on a fixed time grid and preprocess it again" ) _validate_curve_valid_mask( source[mask_key], len(source[f"Y_curve_{split}"]), migration_n_time_points, ) payload = dict(source) for split in ("train", "val", "test"): key = f"Y_curve_{split}" payload[key] = np.asarray(source[key][:, keep_indices], dtype=np.float32) payload["scaler_curve"] = _slice_standard_scaler(source["scaler_curve"], keep_indices) # 这些字段属于已停用的 time-conditioned 路径，不进入旧正演模型训练包。 for key in ( "T_curve_train", "T_curve_val", "T_curve_test", "X_time_train", "X_time_val", "X_time_test", "scaler_time", ): payload.pop(key, None) n_time_points = int(keep_indices.size // 2) meta = dict(source.get("meta", {})) meta.update( { "curve_dim": int(keep_indices.size), "raw_curve_dim": curve_dim, "dropped_legacy_slope": True, "migrated_from_processed": str(input_path), "curve_layout": { "n_time_points": n_time_points, "parts": [ {"name": "log_pressure", "start": 0, "end": n_time_points}, {"name": "log_derivative", "start": n_time_points, "end": 2 * n_time_points}, ], }, } ) meta.pop("time_feature_dim", None) meta.pop("curve_time_source", None) payload["meta"] = meta output_path.parent.mkdir(parents=True, exist_ok=True) joblib.dump(payload, output_path) print(f"Processed migration complete, saved to: {output_path}") print( f"train={len(payload['Y_curve_train'])}, val={len(payload['Y_curve_val'])}, " f"test={len(payload['Y_curve_test'])}, curve_dim={curve_dim}->{keep_indices.size}" ) def preprocess_dataset( input_path: Path, output_path: Path, test_size: float = 0.15, val_size: float = 0.15, random_seed: int = 42, use_param_feature_transform: bool = True, ) -> None: """将原始 HDF5 样本整理为训练脚本可直接加载的 joblib 数据包。 处理步骤包括：读取 params/schedule/curve 等核心数组，保留可选元数据，构造参数 特征变换，划分 train/val/test，并只在训练集上拟合 StandardScaler。输出 payload 同时保存 scaler 和 meta，确保后续评估可以恢复原始尺度并正确拆分压力和导数。 """ if not input_path.exists(): raise FileNotFoundError(f"Input file not found: {input_path}") with h5py.File(input_path, "r") as f: required_keys = ["params", "schedule", "curve"] for key in required_keys: if key not in f: raise KeyError(f"HDF5 dataset is missing required key: {key}") x_params = _ensure_2d("params", f["params"][:]) x_schedule = _ensure_2d("schedule", f["schedule"][:]) y_curve = _ensure_2d("curve", f["curve"][:]) curve_time = _read_optional_numeric_dataset(f, "curve_time") curve_valid_mask = _read_optional_numeric_dataset(f, "curve_valid_mask") # 可选元数据会保留下来，供后续分析和排序验证使用； # 基础正演训练循环不依赖这些字段。 param_names = _decode_attr_list(f.attrs.get("param_names")) schedule_meta_names = _decode_attr_list(f.attrs.get("schedule_meta_names")) source_name_vocab = _decode_attr_list(f.attrs.get("source_name_vocab")) if param_names is None: param_names = ["k", "skin", "wellboreC", "phi", "h", "Ct", "Cf"] solver_type = None categorical_values: dict[str, list[float]] = {} if "solverType" in param_names: solver_col = param_names.index("solverType") solver_type = np.asarray(x_params[:, solver_col], dtype=np.float64) categories = np.unique(solver_type) if not np.all(np.isfinite(categories)): raise ValueError("solverType contains non-finite values") categorical_values["solverType"] = categories.tolist() schedule_meta = _ensure_2d("schedule_meta", f["schedule_meta"][:]) if "schedule_meta" in f else None family_name = _read_optional_string_dataset(f, "family_name") source_id = _read_optional_numeric_dataset(f, "source_id") source_name = _read_optional_string_dataset(f, "source_name") source_row = _read_optional_numeric_dataset(f, "source_row") extra_numeric_names = [ "group_id", "is_anchor", "neighbor_objective", "neighbor_objective_p", "neighbor_objective_d", "neighbor_span_frac", "section_index", ] extra_string_names = ["timeQ_json", "q_json"] extra_numeric = {name: _read_optional_numeric_dataset(f, name) for name in extra_numeric_names} extra_string = {name: _read_optional_string_dataset(f, name) for name in extra_string_names} n = len(x_params) if len(x_schedule) != n or len(y_curve) != n: raise ValueError("params / schedule / curve row counts do not match") raw_curve_dim = int(y_curve.shape[1]) if curve_time is not None: curve_time_shape = np.asarray(curve_time).shape if len(curve_time_shape) != 2 or curve_time_shape[0] != n: raise ValueError( f"curve_time must be a 2D array with {n} rows, got shape={curve_time_shape}" ) n_time_points = int(curve_time_shape[1]) if raw_curve_dim == 3 * n_time_points: # 兼容已有的压力/导数/斜率三通道 HDF5；训练时丢弃斜率通道。 y_curve = y_curve[:, : 2 * n_time_points] elif raw_curve_dim != 2 * n_time_points: raise ValueError( "curve dimension does not match curve_time: " f"curve_dim={raw_curve_dim}, n_time_points={n_time_points}; " "expected 2*N or legacy 3*N" ) elif raw_curve_dim % 3 == 0: n_time_points = raw_curve_dim // 3 y_curve = y_curve[:, : 2 * n_time_points] elif raw_curve_dim % 2 == 0: n_time_points = raw_curve_dim // 2 else: raise ValueError( f"curve dimension {raw_curve_dim} cannot be interpreted as 2*N or legacy 3*N" ) curve_dim = int(y_curve.shape[1]) curve_time_mode = "missing" prediction_curve_time = None if curve_time is not None: curve_time = _validate_curve_time(curve_time, n, n_time_points) if np.allclose(curve_time, curve_time[0:1], rtol=1e-6, atol=1e-10): curve_time_mode = "fixed" prediction_curve_time = curve_time[0].copy() else: curve_time_mode = "per_sample" if curve_time_mode != "fixed": raise ValueError( "forward-surrogate preprocessing requires one fixed curve_time grid; " f"got curve_time_mode={curve_time_mode!r}" ) if curve_valid_mask is None: raise ValueError( "fixed-time preprocessing requires curve_valid_mask; regenerate the H5 dataset" ) curve_valid_mask = _validate_curve_valid_mask( curve_valid_mask, n, n_time_points, ) curve_valid_mask_mode = "explicit" curve_mask = np.concatenate([curve_valid_mask, curve_valid_mask], axis=1) param_feature_transform = build_param_feature_transform( param_names=param_names, categorical_values=categorical_values, enabled=bool(use_param_feature_transform), ) # 物理参数在 StandardScaler 前先做特征变换： # 例如 k、wellboreC、h 使用 log10，skin 使用 asinh。 x_params_features = transform_param_features(x_params, param_feature_transform) _validate_optional_length("schedule_meta", schedule_meta, n) _validate_optional_length("family_name", family_name, n) _validate_optional_length("source_id", source_id, n) _validate_optional_length("source_name", source_name, n) _validate_optional_length("source_row", source_row, n) for name, arr in extra_numeric.items(): _validate_optional_length(name, arr, n) for name, arr in extra_string.items(): _validate_optional_length(name, arr, n) # 同一group通常对应相同流量制度在不同solverType下的样本，必须整体进入同一划分。 group_id = extra_numeric.get("group_id") idx_train, idx_val, idx_test, split_strategy = _split_sample_indices( n_samples=n, test_size=test_size, val_size=val_size, random_seed=random_seed, group_id=group_id, solver_type=solver_type, ) x_params_train = x_params_features[idx_train] x_schedule_train = x_schedule[idx_train] y_curve_train = y_curve[idx_train] curve_valid_mask_train = curve_valid_mask[idx_train] curve_mask_train = curve_mask[idx_train] x_params_val = x_params_features[idx_val] x_schedule_val = x_schedule[idx_val] y_curve_val = y_curve[idx_val] curve_valid_mask_val = curve_valid_mask[idx_val] curve_mask_val = curve_mask[idx_val] x_params_test = x_params_features[idx_test] x_schedule_test = x_schedule[idx_test] y_curve_test = y_curve[idx_test] curve_valid_mask_test = curve_valid_mask[idx_test] curve_mask_test = curve_mask[idx_test] schedule_meta_train = schedule_meta[idx_train] if schedule_meta is not None else None schedule_meta_val = schedule_meta[idx_val] if schedule_meta is not None else None schedule_meta_test = schedule_meta[idx_test] if schedule_meta is not None else None family_name_train = family_name[idx_train] if family_name is not None else None family_name_val = family_name[idx_val] if family_name is not None else None family_name_test = family_name[idx_test] if family_name is not None else None source_id_train = source_id[idx_train] if source_id is not None else None source_id_val = source_id[idx_val] if source_id is not None else None source_id_test = source_id[idx_test] if source_id is not None else None source_name_train = source_name[idx_train] if source_name is not None else None source_name_val = source_name[idx_val] if source_name is not None else None source_name_test = source_name[idx_test] if source_name is not None else None source_row_train = source_row[idx_train] if source_row is not None else None source_row_val = source_row[idx_val] if source_row is not None else None source_row_test = source_row[idx_test] if source_row is not None else None extra_numeric_split = { name: ( arr[idx_train] if arr is not None else None, arr[idx_val] if arr is not None else None, arr[idx_test] if arr is not None else None, ) for name, arr in extra_numeric.items() } extra_string_split = { name: ( arr[idx_train] if arr is not None else None, arr[idx_val] if arr is not None else None, arr[idx_test] if arr is not None else None, ) for name, arr in extra_string.items() } categorical_indices = np.asarray( param_feature_transform.get("categorical_feature_indices", []), dtype=np.int64, ) scaled_param_indices = np.setdiff1d( np.arange(x_params_features.shape[1], dtype=np.int64), categorical_indices, ) scaler_params = SelectiveStandardScaler(scaled_param_indices) scaler_schedule = StandardScaler() scaler_curve = MaskedStandardScaler() valid_count_train = curve_valid_mask_train.sum(axis=0).astype(np.int64) if prediction_curve_time is not None and np.any(valid_count_train == 0): missing_columns = np.flatnonzero(valid_count_train == 0).tolist() raise ValueError( "fixed curve_time has training columns without physical coverage: " f"{missing_columns}" ) # 只在训练集上拟合 scaler，避免验证/测试信息泄漏。 x_params_train_s = scaler_params.fit_transform(x_params_train) x_schedule_train_s = scaler_schedule.fit_transform(x_schedule_train) scaler_curve.fit(y_curve_train, curve_mask_train) y_curve_train_s = scaler_curve.transform(y_curve_train) y_curve_train_s[~curve_mask_train] = 0.0 x_params_val_s = scaler_params.transform(x_params_val) x_schedule_val_s = scaler_schedule.transform(x_schedule_val) y_curve_val_s = scaler_curve.transform(y_curve_val) y_curve_val_s[~curve_mask_val] = 0.0 x_params_test_s = scaler_params.transform(x_params_test) x_schedule_test_s = scaler_schedule.transform(x_schedule_test) y_curve_test_s = scaler_curve.transform(y_curve_test) y_curve_test_s[~curve_mask_test] = 0.0 output_path.parent.mkdir(parents=True, exist_ok=True) # 保存 curve_layout，训练和评估脚本就能按元数据切分拼接曲线， # 避免在各处硬编码索引范围。 meta = { "n_total": int(n), "n_train": int(len(idx_train)), "n_val": int(len(idx_val)), "n_test": int(len(idx_test)), "param_dim": int(x_params_features.shape[1]), "raw_param_dim": int(x_params.shape[1]), "schedule_dim": int(x_schedule.shape[1]), "curve_dim": curve_dim, "raw_curve_dim": raw_curve_dim, "dropped_legacy_slope": bool(raw_curve_dim != curve_dim), "curve_time_mode": curve_time_mode, "curve_valid_mask_mode": curve_valid_mask_mode, "curve_valid_fraction": { "train": float(np.mean(curve_valid_mask_train)), "val": float(np.mean(curve_valid_mask_val)), "test": float(np.mean(curve_valid_mask_test)), }, "curve_valid_count_train": valid_count_train.astype(int).tolist(), "input_h5": str(input_path), "param_names": param_names, "param_feature_transform": param_feature_transform, "param_feature_names": list(param_feature_transform.get("feature_names", [])), "param_scaled_indices": scaled_param_indices.tolist(), "param_unscaled_indices": categorical_indices.tolist(), "split_strategy": split_strategy, "solver_type_counts": { "train": _value_counts(solver_type, idx_train), "val": _value_counts(solver_type, idx_val), "test": _value_counts(solver_type, idx_test), }, "group_counts": { "train": int(np.unique(group_id[idx_train]).size) if group_id is not None else 0, "val": int(np.unique(group_id[idx_val]).size) if group_id is not None else 0, "test": int(np.unique(group_id[idx_test]).size) if group_id is not None else 0, }, "schedule_meta_names": schedule_meta_names, "source_name_vocab": source_name_vocab, "extra_numeric_fields": [name for name, arr in extra_numeric.items() if arr is not None], "extra_string_fields": [name for name, arr in extra_string.items() if arr is not None], "curve_layout": { "n_time_points": int(n_time_points), "parts": [ {"name": "log_pressure", "start": 0, "end": int(n_time_points)}, {"name": "log_derivative", "start": int(n_time_points), "end": int(2 * n_time_points)}, ], }, } if prediction_curve_time is not None: meta["prediction_curve_time"] = prediction_curve_time.tolist() payload = { "X_params_train": x_params_train_s.astype(np.float32), "X_schedule_train": x_schedule_train_s.astype(np.float32), "Y_curve_train": y_curve_train_s.astype(np.float32), "X_params_val": x_params_val_s.astype(np.float32), "X_schedule_val": x_schedule_val_s.astype(np.float32), "Y_curve_val": y_curve_val_s.astype(np.float32), "X_params_test": x_params_test_s.astype(np.float32), "X_schedule_test": x_schedule_test_s.astype(np.float32), "Y_curve_test": y_curve_test_s.astype(np.float32), "curve_valid_mask_train": curve_valid_mask_train.astype(np.uint8), "curve_valid_mask_val": curve_valid_mask_val.astype(np.uint8), "curve_valid_mask_test": curve_valid_mask_test.astype(np.uint8), "scaler_params": scaler_params, "scaler_schedule": scaler_schedule, "scaler_curve": scaler_curve, "meta": meta, } if schedule_meta is not None: payload["schedule_meta_train"] = np.asarray(schedule_meta_train, dtype=np.float32) payload["schedule_meta_val"] = np.asarray(schedule_meta_val, dtype=np.float32) payload["schedule_meta_test"] = np.asarray(schedule_meta_test, dtype=np.float32) if family_name is not None: payload["family_name_train"] = np.asarray(family_name_train, dtype=object) payload["family_name_val"] = np.asarray(family_name_val, dtype=object) payload["family_name_test"] = np.asarray(family_name_test, dtype=object) if source_id is not None: payload["source_id_train"] = np.asarray(source_id_train) payload["source_id_val"] = np.asarray(source_id_val) payload["source_id_test"] = np.asarray(source_id_test) if source_name is not None: payload["source_name_train"] = np.asarray(source_name_train, dtype=object) payload["source_name_val"] = np.asarray(source_name_val, dtype=object) payload["source_name_test"] = np.asarray(source_name_test, dtype=object) if source_row is not None: payload["source_row_train"] = np.asarray(source_row_train) payload["source_row_val"] = np.asarray(source_row_val) payload["source_row_test"] = np.asarray(source_row_test) for name, (train_arr, val_arr, test_arr) in extra_numeric_split.items(): if train_arr is None: continue payload[f"{name}_train"] = np.asarray(train_arr) payload[f"{name}_val"] = np.asarray(val_arr) payload[f"{name}_test"] = np.asarray(test_arr) for name, (train_arr, val_arr, test_arr) in extra_string_split.items(): if train_arr is None: continue payload[f"{name}_train"] = np.asarray(train_arr, dtype=object) payload[f"{name}_val"] = np.asarray(val_arr, dtype=object) payload[f"{name}_test"] = np.asarray(test_arr, dtype=object) joblib.dump(payload, output_path) print(f"Preprocess complete, saved to: {output_path}") print( f"train={len(idx_train)}, val={len(idx_val)}, test={len(idx_test)}, " f"param_dim={x_params_features.shape[1]}, schedule_dim={x_schedule.shape[1]}, " f"curve_dim={y_curve.shape[1]}" ) print( f"split_strategy={split_strategy}, " f"solver_type_counts={meta['solver_type_counts']}, " f"group_counts={meta['group_counts']}, " f"curve_time_mode={curve_time_mode}" ) if schedule_meta is not None: print(f"schedule_meta_dim={schedule_meta.shape[1]}, family_name_saved={family_name is not None}") if source_name is not None: print(f"source_name_saved=True, source_vocab={source_name_vocab}") if meta["extra_numeric_fields"] or meta["extra_string_fields"]: print( "extra_fields_saved=" f"{meta['extra_numeric_fields'] + meta['extra_string_fields']}" )