多相态数据预处理

ML/nmWTAI-ML/src/data/param_features.py

@@ -21,7 +21,7 @@ import numpy as np
 
 # 部分物理参数跨越多个数量级，因此先做特征变换，再交给模型学习。
 DEFAULT_PARAM_NAMES = ["k", "skin", "wellboreC", "phi", "h", "Ct", "Cf"]
-DEFAULT_LOG_PARAM_NAMES = {"k", "wellboreC", "h", "Ct"}
+DEFAULT_LOG_PARAM_NAMES = {"k", "wellboreC", "h", "Ct", "Cf"}
 DEFAULT_ASINH_PARAM_NAMES = {"skin"}
 DEFAULT_COMPOSITE_FEATURES = [
     "log10_kh",
@@ -46,6 +46,7 @@ def build_param_feature_transform(
     param_names: list[str] | None = None,
     log_param_names: set[str] | None = None,
     asinh_param_names: set[str] | None = None,
+    categorical_values: dict[str, list[float]] | None = None,
     enabled: bool = True,
     include_composite_features: bool = True,
 ) -> dict[str, Any]:
@@ -53,12 +54,26 @@ def build_param_feature_transform(
     names = list(param_names or DEFAULT_PARAM_NAMES)
     log_names = set(DEFAULT_LOG_PARAM_NAMES if log_param_names is None else log_param_names)
     asinh_names = set(DEFAULT_ASINH_PARAM_NAMES if asinh_param_names is None else asinh_param_names)
+    categories = {
+        str(name): [float(value) for value in values]
+        for name, values in (categorical_values or {}).items()
+    }
 
     feature_names: list[str] = []
     transforms: dict[str, str] = {}
+    feature_slices: dict[str, list[int]] = {}
+    categorical_feature_indices: list[int] = []
     for name in names:
+        start = len(feature_names)
+        category_values = categories.get(name, [])
         # 单参数特征名称直接记录变换方式，便于训练后反查每一列的物理含义。
-        if enabled and name in log_names:
+        if enabled and category_values:
+            transforms[name] = "one_hot"
+            for value in category_values:
+                suffix = str(int(value)) if float(value).is_integer() else f"{value:g}"
+                feature_names.append(f"is_T{suffix}" if name == "solverType" else f"{name}_{suffix}")
+            categorical_feature_indices.extend(range(start, len(feature_names)))
+        elif enabled and name in log_names:
             transforms[name] = "log10"
             feature_names.append(f"log10_{name}")
         elif enabled and name in asinh_names:
@@ -67,6 +82,7 @@ def build_param_feature_transform(
         else:
             transforms[name] = "identity"
             feature_names.append(name)
+        feature_slices[name] = [start, len(feature_names)]
 
     composite_features = list(DEFAULT_COMPOSITE_FEATURES) if (enabled and include_composite_features) else []
     # 复合特征编码常见试井组合量，例如 kh、k/phi 和井筒储集相关比值。
@@ -77,6 +93,10 @@ def build_param_feature_transform(
         "param_names": names,
         "feature_names": feature_names,
         "transforms": transforms,
+        "feature_slices": feature_slices,
+        "categorical_values": categories,
+        "categorical_feature_indices": categorical_feature_indices,
+        "base_feature_dim": len(feature_names) - len(composite_features),
         "log_param_names": sorted(log_names),
         "asinh_param_names": sorted(asinh_names),
         "composite_features": composite_features,
@@ -94,6 +114,33 @@ def param_feature_transform_from_meta(meta: dict[str, Any] | None) -> dict[str,
     return dict(transform)
 
 
+def build_raw_param_vector(
+    params: Any,
+    transform: dict[str, Any] | None,
+    solver_type: int | float | None = None,
+) -> np.ndarray:
+    """Build raw model input columns in the exact order saved during preprocessing."""
+    if transform is None:
+        # Preserve the input layout used by legacy processed datasets.
+        names = ["k", "skin", "wellboreC", "phi", "h", "Cf"]
+    else:
+        names = list(transform.get("param_names") or DEFAULT_PARAM_NAMES)
+
+    values: list[float] = []
+    for name in names:
+        if name == "solverType":
+            if solver_type is None:
+                raise ValueError("solverType is required by this processed dataset")
+            values.append(float(solver_type))
+            continue
+
+        if not hasattr(params, name):
+            raise ValueError(f"parameter object has no field required by the model: {name}")
+        values.append(float(getattr(params, name)))
+
+    return np.asarray(values, dtype=np.float32).reshape(1, -1)
+
+
 def transform_param_features(
     params: np.ndarray,
     transform: dict[str, Any] | None,
@@ -112,21 +159,35 @@ def transform_param_features(
         raise ValueError(f"param feature transform expects {len(names)} columns, got {x.shape[1]}")
 
     raw = x.astype(np.float64, copy=True)
-    out = raw.copy()
+    columns = []
     log_eps = float(transform.get("log_eps", 1.0e-30))
     for col, name in enumerate(names):
         mode = str(transforms.get(name, "identity")).lower()
         if mode == "log10":
             # 对跨数量级参数取 log10，减少大数值范围对 StandardScaler 和网络的压力。
-            out[:, col] = np.log10(np.maximum(out[:, col], log_eps))
+            columns.append(np.log10(np.maximum(raw[:, col], log_eps)).reshape(-1, 1))
         elif mode == "asinh":
             # skin 可能为负，asinh 保留符号且在大值区域近似对数。
-            out[:, col] = np.arcsinh(out[:, col])
+            columns.append(np.arcsinh(raw[:, col]).reshape(-1, 1))
         elif mode == "identity":
-            continue
+            columns.append(raw[:, col].reshape(-1, 1))
+        elif mode == "one_hot":
+            category_values = np.asarray(
+                (transform.get("categorical_values") or {}).get(name, []),
+                dtype=np.float64,
+            )
+            if category_values.size == 0:
+                raise ValueError(f"one-hot parameter {name} has no configured categories")
+            encoded = np.isclose(raw[:, col : col + 1], category_values.reshape(1, -1))
+            if not np.all(np.any(encoded, axis=1)):
+                unknown = np.unique(raw[~np.any(encoded, axis=1), col]).tolist()
+                raise ValueError(f"Unknown categorical values for {name}: {unknown}")
+            columns.append(encoded.astype(np.float64))
         else:
             raise ValueError(f"Unknown transform mode for {name}: {mode}")
 
+    out = np.concatenate(columns, axis=1) if columns else np.empty((raw.shape[0], 0), dtype=np.float64)
+
     composite_features = list(transform.get("composite_features") or [])
     if composite_features:
         name_to_col = {name: idx for idx, name in enumerate(names)}
@@ -174,18 +235,32 @@ def inverse_transform_param_features(
 
     names = list(transform.get("param_names") or DEFAULT_PARAM_NAMES)
     transforms = dict(transform.get("transforms") or {})
+    feature_slices = dict(transform.get("feature_slices") or {})
+    if not feature_slices:
         if x.shape[1] < len(names):
             raise ValueError(f"param inverse transform expects at least {len(names)} columns, got {x.shape[1]}")
+        feature_slices = {name: [idx, idx + 1] for idx, name in enumerate(names)}
 
-    out = x[:, : len(names)].astype(np.float64, copy=True)
+    out = np.empty((x.shape[0], len(names)), dtype=np.float64)
     for col, name in enumerate(names):
+        start, end = map(int, feature_slices[name])
+        if end > x.shape[1]:
+            raise ValueError(f"param inverse transform slice for {name} exceeds feature dimension")
         mode = str(transforms.get(name, "identity")).lower()
-        if mode == "log10":
+        if mode == "one_hot":
-            out[:, col] = 10.0 ** out[:, col]
+            category_values = np.asarray(
+                (transform.get("categorical_values") or {}).get(name, []),
+                dtype=np.float64,
+            )
+            if category_values.size != end - start:
+                raise ValueError(f"one-hot category count mismatch for {name}")
+            out[:, col] = category_values[np.argmax(x[:, start:end], axis=1)]
+        elif mode == "log10":
+            out[:, col] = 10.0 ** x[:, start]
         elif mode == "asinh":
-            out[:, col] = np.sinh(out[:, col])
+            out[:, col] = np.sinh(x[:, start])
         elif mode == "identity":
-            continue
+            out[:, col] = x[:, start]
         else:
             raise ValueError(f"Unknown transform mode for {name}: {mode}")
     return out.astype(np.float32)

ML/nmWTAI-ML/src/data/preprocess.py

@@ -18,12 +18,137 @@ from pathlib import Path
 import h5py
 import joblib
 import numpy as np
-from sklearn.model_selection import train_test_split
+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)
+
+
+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:
@@ -189,6 +314,16 @@ def preprocess_dataset(
         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")
@@ -227,6 +362,7 @@ def preprocess_dataset(
 
     param_feature_transform = build_param_feature_transform(
         param_names=param_names,
+        categorical_values=categorical_values,
         enabled=bool(use_param_feature_transform),
     )
     # 物理参数在 StandardScaler 前先做特征变换：
@@ -243,21 +379,15 @@ def preprocess_dataset(
     for name, arr in extra_string.items():
         _validate_optional_length(name, arr, n)
 
-    idx = np.arange(n)
+    # 同一group通常对应相同流量制度在不同solverType下的样本，必须整体进入同一划分。
-    # 先划出测试集，再从剩余样本中划分训练集和验证集。
+    group_id = extra_numeric.get("group_id")
-    idx_train_val, idx_test = train_test_split(
+    idx_train, idx_val, idx_test, split_strategy = _split_sample_indices(
-        idx,
+        n_samples=n,
         test_size=test_size,
-        random_state=random_seed,
+        val_size=val_size,
-        shuffle=True,
+        random_seed=random_seed,
-    )
+        group_id=group_id,
-
+        solver_type=solver_type,
-    val_ratio_in_train_val = val_size / (1.0 - test_size)
-    idx_train, idx_val = train_test_split(
-        idx_train_val,
-        test_size=val_ratio_in_train_val,
-        random_state=random_seed,
-        shuffle=True,
     )
 
     x_params_train = x_params_features[idx_train]
@@ -309,7 +439,15 @@ def preprocess_dataset(
         for name, arr in extra_string.items()
     }
 
-    scaler_params = StandardScaler()
+    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 = StandardScaler()
 
@@ -344,6 +482,20 @@ def preprocess_dataset(
         "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],
@@ -414,6 +566,11 @@ def preprocess_dataset(
         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']}"
+    )
     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: