|
|
from __future__ import annotations
|
|
|
|
|
|
import json
|
|
|
import multiprocessing as mp
|
|
|
import os
|
|
|
import shutil
|
|
|
import struct
|
|
|
import subprocess
|
|
|
from collections import deque
|
|
|
from concurrent.futures import ProcessPoolExecutor
|
|
|
from datetime import datetime
|
|
|
from pathlib import Path
|
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
|
|
import h5py
|
|
|
import numpy as np
|
|
|
try:
|
|
|
from tqdm import tqdm
|
|
|
except ImportError:
|
|
|
def tqdm(iterable=None, **kwargs):
|
|
|
if iterable is None:
|
|
|
class _NoopTqdm:
|
|
|
def __init__(self, total=None, **_):
|
|
|
self.total = total
|
|
|
|
|
|
def update(self, _=1):
|
|
|
pass
|
|
|
|
|
|
def close(self):
|
|
|
pass
|
|
|
|
|
|
return _NoopTqdm(**kwargs)
|
|
|
return iterable
|
|
|
|
|
|
from src.common.config import Config
|
|
|
from src.common.experiment_paths import normalize_tag
|
|
|
from src.data.curve_processing import clean_curve_for_dataset, is_valid_curve, resample_curve_to_features_with_time
|
|
|
from src.data.params import Params, Schedule, generate_params_dataset, write_params_bin
|
|
|
from src.data.schedule_encoding import encode_schedule_to_timegrid
|
|
|
|
|
|
_WORKER_CFG_PATH: Optional[str] = None
|
|
|
_WORKDIR: Optional[Path] = None
|
|
|
_PARAMS_BIN: Optional[Path] = None
|
|
|
_RESULT_BIN: Optional[Path] = None
|
|
|
_SUBPROC_ENV: Optional[dict] = None
|
|
|
|
|
|
|
|
|
"""生成正演代理模型训练使用的 HDF5 样本。
|
|
|
|
|
|
每个样本由“物理参数采样 + 流量制度采样 + C++ 数值求解器正演”生成，
|
|
|
随后将求解器输出的双对数曲线转换为固定长度的模型输入和监督目标。
|
|
|
"""
|
|
|
|
|
|
|
|
|
def _seed32(x: int) -> int:
|
|
|
return int(x) & 0xFFFFFFFF
|
|
|
|
|
|
|
|
|
def _read_result_bin(result_bin_path: Path) -> Optional[Dict[str, Any]]:
|
|
|
"""读取 runner.exe 输出的紧凑二进制结果文件。"""
|
|
|
try:
|
|
|
with open(result_bin_path, "rb") as f:
|
|
|
magic, version = struct.unpack("<II", f.read(8))
|
|
|
expected_magic = ord("R") | (ord("S") << 8) | (ord("B") << 16) | (ord("1") << 24)
|
|
|
if magic != expected_magic or version != 1:
|
|
|
return None
|
|
|
|
|
|
nWells, nSteps = struct.unpack("<II", f.read(8))
|
|
|
t = np.fromfile(f, dtype="<f8", count=nSteps).astype(np.float64).tolist()
|
|
|
pw = [np.fromfile(f, dtype="<f8", count=nSteps).astype(np.float64).tolist() for _ in range(nWells)]
|
|
|
|
|
|
loglog = []
|
|
|
for _ in range(nWells):
|
|
|
(nLogLog,) = struct.unpack("<I", f.read(4))
|
|
|
loglog_t = np.fromfile(f, dtype="<f8", count=nLogLog).astype(np.float64).tolist()
|
|
|
loglog_p = np.fromfile(f, dtype="<f8", count=nLogLog).astype(np.float64).tolist()
|
|
|
loglog_deriv = np.fromfile(f, dtype="<f8", count=nLogLog).astype(np.float64).tolist()
|
|
|
loglog.append({"t": loglog_t, "p": loglog_p, "deriv": loglog_deriv})
|
|
|
|
|
|
return {"nWells": nWells, "nSteps": nSteps, "t": t, "pw": pw, "loglog": loglog}
|
|
|
except Exception:
|
|
|
return None
|
|
|
|
|
|
|
|
|
def _pick_mixture(rng, items):
|
|
|
probs = np.array([float(it.get("prob", 0.0)) for it in items], dtype=np.float64)
|
|
|
s = float(np.sum(probs))
|
|
|
if s <= 0:
|
|
|
return items[int(rng.randint(0, len(items)))]
|
|
|
probs = probs / s
|
|
|
|
|
|
u = rng.rand()
|
|
|
c = 0.0
|
|
|
for it, p in zip(items, probs):
|
|
|
c += float(p)
|
|
|
if u <= c:
|
|
|
return it
|
|
|
return items[-1]
|
|
|
|
|
|
|
|
|
def _normalize_durations_to_total(dt: np.ndarray, total: float, min_dt: float) -> np.ndarray:
|
|
|
dt = np.maximum(dt.astype(np.float64), float(min_dt))
|
|
|
s = float(np.sum(dt))
|
|
|
total = max(float(total), float(min_dt) * float(len(dt)))
|
|
|
if s <= 0:
|
|
|
return np.full_like(dt, total / len(dt))
|
|
|
dt = dt * (total / s)
|
|
|
dt = np.maximum(dt, float(min_dt))
|
|
|
diff = total - float(np.sum(dt))
|
|
|
dt[-1] += diff
|
|
|
return np.maximum(dt, float(min_dt))
|
|
|
|
|
|
|
|
|
SCHEDULE_META_NAMES = [
|
|
|
"family_id",
|
|
|
"section_index",
|
|
|
"n_sections",
|
|
|
"n_prod",
|
|
|
"prod_total_time",
|
|
|
"shutin_dt",
|
|
|
"q_first_prod",
|
|
|
"q_last_prod",
|
|
|
"q_mean_prod",
|
|
|
"q_std_prod",
|
|
|
"q_min_prod",
|
|
|
"q_max_prod",
|
|
|
"q_end_start_ratio",
|
|
|
"q_range_ratio",
|
|
|
"max_step_ratio",
|
|
|
"mean_step_ratio",
|
|
|
"trend_slope",
|
|
|
"is_nonincreasing",
|
|
|
"is_strong_decline",
|
|
|
]
|
|
|
|
|
|
|
|
|
def _family_id_map(cfg: Config) -> dict[str, int]:
|
|
|
mode = str(cfg.raw["schedule"]["generation_mode"]).lower()
|
|
|
if mode == "family_random":
|
|
|
families = cfg.raw["schedule"]["family_random"]["families"]
|
|
|
mapping = {str(item.get("name", f"family_{i}")).lower(): i for i, item in enumerate(families)}
|
|
|
mapping.setdefault("unknown", -1)
|
|
|
return mapping
|
|
|
return {
|
|
|
"fixed_case": 0,
|
|
|
"case_neighborhood": 1,
|
|
|
"unknown": -1,
|
|
|
}
|
|
|
|
|
|
|
|
|
def build_schedule_metadata(
|
|
|
cfg: Config,
|
|
|
timeQ: List[float],
|
|
|
q: List[float],
|
|
|
family_name: str,
|
|
|
section_index: int,
|
|
|
) -> tuple[np.ndarray, str]:
|
|
|
"""构造描述流量制度形态的分析元数据。"""
|
|
|
dt = np.asarray(timeQ, dtype=np.float64).reshape(-1)
|
|
|
qq = np.asarray(q, dtype=np.float64).reshape(-1)
|
|
|
n_sections = int(len(dt))
|
|
|
q_thr = float((cfg.raw["schedule"].get("canonicalize_for_model", {}) or {}).get("q_thr", 1e-6))
|
|
|
|
|
|
has_shutin = bool(n_sections > 0 and qq[-1] <= q_thr)
|
|
|
n_prod = int(max(n_sections - 1, 0) if has_shutin else n_sections)
|
|
|
prod_dt = dt[:n_prod] if n_prod > 0 else np.zeros((0,), dtype=np.float64)
|
|
|
prod_q = qq[:n_prod] if n_prod > 0 else np.zeros((0,), dtype=np.float64)
|
|
|
shutin_dt = float(dt[-1]) if has_shutin else 0.0
|
|
|
q_eps = 1.0e-12
|
|
|
|
|
|
if n_prod > 0:
|
|
|
q_first = float(prod_q[0])
|
|
|
q_last = float(prod_q[-1])
|
|
|
q_min = float(np.min(prod_q))
|
|
|
q_max = float(np.max(prod_q))
|
|
|
q_end_start_ratio = q_last / max(abs(q_first), q_eps)
|
|
|
q_range_ratio = (q_max - q_min) / max(abs(q_max), q_eps)
|
|
|
else:
|
|
|
q_first = q_last = q_min = q_max = 0.0
|
|
|
q_end_start_ratio = 0.0
|
|
|
q_range_ratio = 0.0
|
|
|
|
|
|
if n_prod > 1:
|
|
|
prev_q = prod_q[:-1]
|
|
|
next_q = prod_q[1:]
|
|
|
step_ratio = np.maximum(prev_q, next_q) / np.maximum(np.minimum(prev_q, next_q), q_eps)
|
|
|
max_step_ratio = float(np.max(step_ratio))
|
|
|
mean_step_ratio = float(np.mean(step_ratio))
|
|
|
is_nonincreasing = bool(np.all(next_q <= prev_q * (1.0 + 1.0e-9)))
|
|
|
has_drop = bool(np.any(next_q < prev_q * (1.0 - 1.0e-9)))
|
|
|
x = np.cumsum(prod_dt) - 0.5 * prod_dt
|
|
|
trend_slope = float(np.polyfit(x, prod_q, deg=1)[0]) if np.ptp(x) > q_eps else 0.0
|
|
|
else:
|
|
|
max_step_ratio = 1.0
|
|
|
mean_step_ratio = 1.0
|
|
|
is_nonincreasing = False
|
|
|
has_drop = False
|
|
|
trend_slope = 0.0
|
|
|
|
|
|
is_strong_decline = bool(is_nonincreasing and has_drop and q_end_start_ratio < 0.70)
|
|
|
|
|
|
family_name = str(family_name).lower()
|
|
|
family_id = int(_family_id_map(cfg).get(family_name, -1))
|
|
|
|
|
|
meta_vec = np.asarray(
|
|
|
[
|
|
|
float(family_id),
|
|
|
float(section_index),
|
|
|
float(n_sections),
|
|
|
float(n_prod),
|
|
|
float(np.sum(prod_dt)) if n_prod > 0 else 0.0,
|
|
|
float(shutin_dt),
|
|
|
q_first,
|
|
|
q_last,
|
|
|
float(np.mean(prod_q)) if n_prod > 0 else 0.0,
|
|
|
float(np.std(prod_q)) if n_prod > 0 else 0.0,
|
|
|
q_min,
|
|
|
q_max,
|
|
|
q_end_start_ratio,
|
|
|
q_range_ratio,
|
|
|
max_step_ratio,
|
|
|
mean_step_ratio,
|
|
|
trend_slope,
|
|
|
float(is_nonincreasing),
|
|
|
float(is_strong_decline),
|
|
|
],
|
|
|
dtype=np.float32,
|
|
|
)
|
|
|
return meta_vec, family_name
|
|
|
|
|
|
|
|
|
def _sample_schedule_fixed_case(cfg: Config, rng) -> Tuple[List[float], List[float], Dict[str, Any]]:
|
|
|
sc = cfg.raw["schedule"]["case_schedule"]
|
|
|
return (
|
|
|
list(map(float, sc["timeQ"])),
|
|
|
list(map(float, sc["q"])),
|
|
|
{"family_name": "fixed_case"},
|
|
|
)
|
|
|
|
|
|
|
|
|
def _sample_schedule_case_neighborhood(cfg: Config, rng) -> Tuple[List[float], List[float], Dict[str, Any]]:
|
|
|
base = cfg.raw["schedule"]["case_schedule"]
|
|
|
ncfg = cfg.raw["schedule"]["case_neighborhood"]
|
|
|
|
|
|
base_timeQ = np.asarray(base["timeQ"], dtype=np.float64)
|
|
|
base_q = np.asarray(base["q"], dtype=np.float64)
|
|
|
out_t = base_timeQ.copy()
|
|
|
out_q = base_q.copy()
|
|
|
prod_n = len(base_timeQ) - 1
|
|
|
|
|
|
noise_t = rng.uniform(-float(ncfg["dt_jitter_rel"]), float(ncfg["dt_jitter_rel"]), size=prod_n)
|
|
|
out_t[:prod_n] = np.maximum(base_timeQ[:prod_n] * (1.0 + noise_t), float(ncfg["min_dt"]))
|
|
|
tail_noise = rng.uniform(-float(ncfg["shutin_dt_jitter_rel"]), float(ncfg["shutin_dt_jitter_rel"]))
|
|
|
out_t[-1] = max(base_timeQ[-1] * (1.0 + tail_noise), float(ncfg["min_dt"]))
|
|
|
|
|
|
noise_q = rng.uniform(-float(ncfg["q_jitter_rel"]), float(ncfg["q_jitter_rel"]), size=prod_n)
|
|
|
out_q[:prod_n] = np.clip(base_q[:prod_n] * (1.0 + noise_q), float(ncfg["q_min"]), float(ncfg["q_max"]))
|
|
|
|
|
|
if bool(ncfg["keep_monotonic_prod"]) and prod_n > 1:
|
|
|
out_q[:prod_n] = np.maximum.accumulate(out_q[:prod_n])
|
|
|
if bool(ncfg["keep_last_q_zero"]):
|
|
|
out_q[-1] = 0.0
|
|
|
|
|
|
return out_t.tolist(), out_q.tolist(), {"family_name": "case_neighborhood"}
|
|
|
|
|
|
|
|
|
def _sample_schedule_family_random(cfg: Config, rng) -> Tuple[List[float], List[float], Dict[str, Any]]:
|
|
|
"""随机抽取一种流量制度族，并生成对应的流动段时长和流量。"""
|
|
|
fcfg = cfg.raw["schedule"]["family_random"]
|
|
|
fam = _pick_mixture(rng, fcfg["families"])
|
|
|
fam_name = str(fam.get("name", "inc_tail_shutin")).lower()
|
|
|
|
|
|
n_lo, n_hi = fcfg["n_prod_sections_range"]
|
|
|
n_prod = int(rng.randint(int(n_lo), int(n_hi) + 1))
|
|
|
|
|
|
prod_total_lo, prod_total_hi = fcfg["prod_total_time_range"]
|
|
|
prod_total = float(prod_total_lo + rng.rand() * (prod_total_hi - prod_total_lo))
|
|
|
|
|
|
mu = float(fcfg["duration_lognormal_mu"])
|
|
|
sigma = float(fcfg["duration_lognormal_sigma"])
|
|
|
dt_prod = rng.lognormal(mean=mu, sigma=sigma, size=n_prod).astype(np.float64)
|
|
|
dt_prod = _normalize_durations_to_total(dt_prod, total=prod_total, min_dt=0.05)
|
|
|
|
|
|
q_lo, q_hi = fcfg["q_range"]
|
|
|
q0 = float(q_lo + rng.rand() * (q_hi - q_lo))
|
|
|
max_rel_step = float(fcfg["max_rel_step"])
|
|
|
mult_noise_sigma = float(fcfg["mult_noise_sigma"])
|
|
|
step_jump_lo, step_jump_hi = fcfg["step_jump_rel_range"]
|
|
|
|
|
|
q_prod = np.zeros((n_prod,), dtype=np.float64)
|
|
|
q_prod[0] = q0
|
|
|
|
|
|
if fam_name == "inc_tail_shutin":
|
|
|
for i in range(1, n_prod):
|
|
|
q_prod[i] = q_prod[i - 1] * rng.uniform(1.02, max_rel_step)
|
|
|
q_prod = np.maximum.accumulate(q_prod)
|
|
|
elif fam_name == "dec_tail_shutin":
|
|
|
for i in range(1, n_prod):
|
|
|
q_prod[i] = max(q_prod[i - 1] * rng.uniform(1.0 / max_rel_step, 0.98), q_lo)
|
|
|
q_prod = np.minimum.accumulate(q_prod)
|
|
|
elif fam_name == "strong_dec_tail_shutin":
|
|
|
end_ratio_lo, end_ratio_hi = fcfg.get("strong_dec_end_start_ratio_range", [0.25, 0.65])
|
|
|
q_end = max(q_lo, q_prod[0] * rng.uniform(float(end_ratio_lo), float(end_ratio_hi)))
|
|
|
shape = np.linspace(0.0, 1.0, n_prod)
|
|
|
q_prod = q_prod[0] * (q_end / max(q_prod[0], 1.0e-12)) ** shape
|
|
|
q_prod = np.minimum.accumulate(q_prod)
|
|
|
elif fam_name == "mild_step_tail_shutin":
|
|
|
q_prod[:] = q0
|
|
|
jump_idx = int(rng.randint(1, max(2, n_prod)))
|
|
|
q_prod[jump_idx:] *= float(rng.uniform(step_jump_lo, step_jump_hi))
|
|
|
else:
|
|
|
for i in range(1, n_prod):
|
|
|
rel = rng.uniform(1.0 / max_rel_step, max_rel_step)
|
|
|
rel = 1.0 + 0.15 * (rel - 1.0)
|
|
|
q_prod[i] = q_prod[i - 1] * rel
|
|
|
|
|
|
if mult_noise_sigma > 0:
|
|
|
q_prod *= np.exp(rng.normal(loc=0.0, scale=mult_noise_sigma, size=n_prod))
|
|
|
|
|
|
q_prod = np.clip(q_prod, q_lo, q_hi)
|
|
|
|
|
|
shut_lo, shut_hi = fcfg["shutin_dt_range"]
|
|
|
dt_shut = float(shut_lo + rng.rand() * (shut_hi - shut_lo))
|
|
|
|
|
|
return dt_prod.tolist() + [dt_shut], q_prod.tolist() + [0.0], {"family_name": fam_name}
|
|
|
|
|
|
|
|
|
def sample_schedule_by_mode(cfg: Config, rng) -> Tuple[List[float], List[float], Dict[str, Any]]:
|
|
|
"""根据配置中的 generation_mode 分发到对应的流量制度采样方法。"""
|
|
|
mode = str(cfg.raw["schedule"]["generation_mode"]).lower()
|
|
|
if mode == "fixed_case":
|
|
|
return _sample_schedule_fixed_case(cfg, rng)
|
|
|
if mode == "case_neighborhood":
|
|
|
return _sample_schedule_case_neighborhood(cfg, rng)
|
|
|
if mode == "family_random":
|
|
|
return _sample_schedule_family_random(cfg, rng)
|
|
|
raise ValueError(f"Unknown schedule generation_mode: {mode}")
|
|
|
|
|
|
|
|
|
def _resolve_section_indices(cfg: Config, timeQ, q, rng):
|
|
|
policy = cfg.raw["schedule"]["section_policy"]
|
|
|
mode = str(policy["mode"]).lower()
|
|
|
n = int(len(timeQ))
|
|
|
if mode == "fixed_last":
|
|
|
return [n]
|
|
|
if mode == "fixed_value":
|
|
|
return [int(np.clip(int(policy["fixed_value"]), 1, n))]
|
|
|
if mode == "all_sections":
|
|
|
return list(range(1, n + 1))
|
|
|
if mode == "uniform_one":
|
|
|
return [int(rng.randint(1, n + 1))]
|
|
|
raise ValueError(f"Unknown section_policy.mode: {mode}")
|
|
|
|
|
|
|
|
|
def _init_process_worker(cfg_path: str):
|
|
|
global _WORKER_CFG_PATH, _WORKDIR, _PARAMS_BIN, _RESULT_BIN, _SUBPROC_ENV
|
|
|
|
|
|
_WORKER_CFG_PATH = cfg_path
|
|
|
cfg = Config(cfg_path)
|
|
|
cfg.ensure_dirs()
|
|
|
|
|
|
base_worker_dir = cfg.paths.temp_dir / "parallel_workers"
|
|
|
pid = os.getpid()
|
|
|
_WORKDIR = base_worker_dir / f"proc_{pid}"
|
|
|
_WORKDIR.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
_PARAMS_BIN = _WORKDIR / "params.bin"
|
|
|
_RESULT_BIN = _WORKDIR / "result.bin"
|
|
|
|
|
|
env = os.environ.copy()
|
|
|
env.setdefault("OMP_NUM_THREADS", "1")
|
|
|
env.setdefault("OPENBLAS_NUM_THREADS", "1")
|
|
|
env.setdefault("MKL_NUM_THREADS", "1")
|
|
|
env.setdefault("NUMEXPR_NUM_THREADS", "1")
|
|
|
_SUBPROC_ENV = env
|
|
|
|
|
|
|
|
|
def _worker_simulate_parallel(args):
|
|
|
"""子进程侧仿真：调用求解器并返回一个编码后的有效样本。"""
|
|
|
task_idx, params_dict = args
|
|
|
cfg = Config(_WORKER_CFG_PATH)
|
|
|
|
|
|
try:
|
|
|
p = Params(
|
|
|
k=float(params_dict["k"]),
|
|
|
skin=float(params_dict["skin"]),
|
|
|
wellboreC=float(params_dict["wellboreC"]),
|
|
|
phi=float(params_dict["phi"]),
|
|
|
h=float(params_dict["h"]),
|
|
|
Cf=float(params_dict["Cf"]),
|
|
|
schedule=None,
|
|
|
)
|
|
|
|
|
|
sc = params_dict["schedule"]
|
|
|
sch = Schedule(
|
|
|
sectionIndex=int(sc["sectionIndex"]),
|
|
|
timeQ=list(map(float, sc["timeQ"])),
|
|
|
q=list(map(float, sc["q"])),
|
|
|
)
|
|
|
p.schedule = sch
|
|
|
|
|
|
write_params_bin(str(_PARAMS_BIN), p, cfg=cfg, include_schedule=True)
|
|
|
|
|
|
if _RESULT_BIN.exists():
|
|
|
_RESULT_BIN.unlink()
|
|
|
|
|
|
cmd = [str(cfg.runner_exe), str(cfg.dataset_bin), str(_PARAMS_BIN), str(_RESULT_BIN)]
|
|
|
result = subprocess.run(
|
|
|
cmd,
|
|
|
cwd=str(cfg.runner_exe.parent),
|
|
|
stdout=subprocess.DEVNULL,
|
|
|
stderr=subprocess.PIPE,
|
|
|
timeout=int(cfg.get("dataset_runtime", "runner_timeout_sec", default=240)),
|
|
|
encoding="utf-8",
|
|
|
errors="ignore",
|
|
|
env=_SUBPROC_ENV,
|
|
|
)
|
|
|
|
|
|
if result.returncode != 0 or not _RESULT_BIN.exists():
|
|
|
return task_idx, None, "runner_failed", None
|
|
|
|
|
|
result_dict = _read_result_bin(_RESULT_BIN)
|
|
|
if not result_dict or not result_dict["loglog"] or not result_dict["loglog"][0]:
|
|
|
return task_idx, None, "empty_loglog", None
|
|
|
|
|
|
t = np.asarray(result_dict["loglog"][0]["t"], dtype=np.float64)
|
|
|
p_curve = np.asarray(result_dict["loglog"][0]["p"], dtype=np.float64)
|
|
|
d_curve = np.asarray(result_dict["loglog"][0]["deriv"], dtype=np.float64)
|
|
|
|
|
|
cleaned = clean_curve_for_dataset(cfg, t, p_curve, d_curve)
|
|
|
if cleaned is None:
|
|
|
return task_idx, None, "too_short_after_clean", None
|
|
|
|
|
|
t, p_curve, d_curve = cleaned
|
|
|
ok, reason = is_valid_curve(cfg, t, p_curve, d_curve)
|
|
|
if not ok:
|
|
|
return task_idx, None, reason, None
|
|
|
|
|
|
curve_feat, curve_time = resample_curve_to_features_with_time(cfg, t, p_curve, d_curve)
|
|
|
|
|
|
sec = int(np.clip(int(sch.sectionIndex), 1, max(len(sch.timeQ), 1)))
|
|
|
enc = encode_schedule_to_timegrid(
|
|
|
cfg,
|
|
|
sectionIndex=sec,
|
|
|
timeQ=sch.timeQ,
|
|
|
q=sch.q,
|
|
|
n_sections=len(sch.timeQ),
|
|
|
)
|
|
|
|
|
|
params_vec = np.asarray(
|
|
|
[p.k, p.skin, p.wellboreC, p.phi, p.h, p.Cf],
|
|
|
dtype=np.float32,
|
|
|
)
|
|
|
|
|
|
schedule_parts = [
|
|
|
np.asarray(enc.x_sched, dtype=np.float32).reshape(-1),
|
|
|
np.asarray(enc.x_sec, dtype=np.float32).reshape(-1),
|
|
|
]
|
|
|
|
|
|
if bool(cfg.raw.get("schedule", {}).get("use_metadata_features_for_model", False)):
|
|
|
names = cfg.raw.get("schedule", {}).get("metadata_features_for_model", []) or []
|
|
|
name_to_idx = {name: i for i, name in enumerate(SCHEDULE_META_NAMES)}
|
|
|
meta_src = np.asarray(params_dict["schedule_meta"]["meta_vec"], dtype=np.float32).reshape(-1)
|
|
|
meta_values = []
|
|
|
for name in names:
|
|
|
idx = name_to_idx.get(str(name))
|
|
|
meta_values.append(float(meta_src[idx]) if idx is not None and idx < meta_src.size else 0.0)
|
|
|
schedule_parts.append(np.asarray(meta_values, dtype=np.float32).reshape(-1))
|
|
|
|
|
|
schedule_vec = np.concatenate(
|
|
|
[
|
|
|
*schedule_parts,
|
|
|
],
|
|
|
axis=0,
|
|
|
).astype(np.float32)
|
|
|
|
|
|
sample = {
|
|
|
"group_id": int(params_dict.get("schedule_id", -1)),
|
|
|
"params": params_vec,
|
|
|
"schedule": schedule_vec,
|
|
|
"curve": np.asarray(curve_feat, dtype=np.float32),
|
|
|
"curve_time": np.asarray(curve_time, dtype=np.float32),
|
|
|
"schedule_meta": np.asarray(params_dict["schedule_meta"]["meta_vec"], dtype=np.float32),
|
|
|
"family_name": str(params_dict["schedule_meta"]["family_name"]),
|
|
|
}
|
|
|
return task_idx, sample, "", None
|
|
|
|
|
|
except subprocess.TimeoutExpired:
|
|
|
return task_idx, None, "timeout", None
|
|
|
except Exception as e:
|
|
|
return task_idx, None, f"exception_{type(e).__name__}", None
|
|
|
|
|
|
|
|
|
class HDF5Appender:
|
|
|
"""将生成样本追加写入 HDF5，避免把完整数据集放入内存。"""
|
|
|
|
|
|
def __init__(
|
|
|
self,
|
|
|
filepath: Path,
|
|
|
param_dim: int,
|
|
|
schedule_dim: int,
|
|
|
curve_dim: int,
|
|
|
compression=None,
|
|
|
chunk_rows: int = 2048,
|
|
|
):
|
|
|
self.filepath = Path(filepath)
|
|
|
self.filepath.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
self.param_dim = int(param_dim)
|
|
|
self.schedule_dim = int(schedule_dim)
|
|
|
self.curve_dim = int(curve_dim)
|
|
|
self.time_dim = int(curve_dim) // 3
|
|
|
|
|
|
self.f = h5py.File(self.filepath, "w")
|
|
|
self._n = 0
|
|
|
self.f.attrs["schedule_meta_names"] = np.asarray(SCHEDULE_META_NAMES, dtype="S")
|
|
|
|
|
|
def ds(name, shape_tail, dtype, fillvalue=None):
|
|
|
# 所有数据集都在第一维支持扩展，用于流式写入样本。
|
|
|
return self.f.create_dataset(
|
|
|
name,
|
|
|
shape=(0, *shape_tail),
|
|
|
maxshape=(None, *shape_tail),
|
|
|
chunks=(chunk_rows, *shape_tail),
|
|
|
dtype=dtype,
|
|
|
compression=compression,
|
|
|
fillvalue=fillvalue,
|
|
|
)
|
|
|
|
|
|
self.d_group_id = ds("group_id", (), np.int64, fillvalue=-1)
|
|
|
self.d_params = ds("params", (self.param_dim,), np.float32, fillvalue=np.nan)
|
|
|
self.d_schedule = ds("schedule", (self.schedule_dim,), np.float32, fillvalue=np.nan)
|
|
|
self.d_curve = ds("curve", (self.curve_dim,), np.float32, fillvalue=np.nan)
|
|
|
self.d_curve_time = ds("curve_time", (self.time_dim,), np.float32, fillvalue=np.nan)
|
|
|
self.d_schedule_meta = ds("schedule_meta", (len(SCHEDULE_META_NAMES),), np.float32, fillvalue=np.nan)
|
|
|
self.d_family_name = self.f.create_dataset(
|
|
|
"family_name",
|
|
|
shape=(0,),
|
|
|
maxshape=(None,),
|
|
|
chunks=(chunk_rows,),
|
|
|
dtype=h5py.string_dtype(encoding="utf-8"),
|
|
|
)
|
|
|
|
|
|
@property
|
|
|
def n_samples(self) -> int:
|
|
|
return int(self._n)
|
|
|
|
|
|
def append_samples(self, samples: List[Dict[str, Any]]):
|
|
|
"""追加一批有效样本，并更新样本数量属性。"""
|
|
|
if not samples:
|
|
|
return
|
|
|
|
|
|
B = len(samples)
|
|
|
start, end = self._n, self._n + B
|
|
|
|
|
|
self.d_group_id.resize((end,))
|
|
|
self.d_params.resize((end, self.param_dim))
|
|
|
self.d_schedule.resize((end, self.schedule_dim))
|
|
|
self.d_curve.resize((end, self.curve_dim))
|
|
|
self.d_curve_time.resize((end, self.time_dim))
|
|
|
self.d_schedule_meta.resize((end, len(SCHEDULE_META_NAMES)))
|
|
|
self.d_family_name.resize((end,))
|
|
|
|
|
|
gid = np.full((B,), -1, dtype=np.int64)
|
|
|
params = np.full((B, self.param_dim), np.nan, dtype=np.float32)
|
|
|
schedule = np.full((B, self.schedule_dim), np.nan, dtype=np.float32)
|
|
|
curve = np.full((B, self.curve_dim), np.nan, dtype=np.float32)
|
|
|
curve_time = np.full((B, self.time_dim), np.nan, dtype=np.float32)
|
|
|
schedule_meta = np.full((B, len(SCHEDULE_META_NAMES)), np.nan, dtype=np.float32)
|
|
|
family_name: list[str] = ["" for _ in range(B)]
|
|
|
|
|
|
for i, s in enumerate(samples):
|
|
|
gid[i] = int(s.get("group_id", -1))
|
|
|
params[i] = np.asarray(s["params"], dtype=np.float32).reshape(-1)[: self.param_dim]
|
|
|
schedule[i] = np.asarray(s["schedule"], dtype=np.float32).reshape(-1)[: self.schedule_dim]
|
|
|
curve[i] = np.asarray(s["curve"], dtype=np.float32).reshape(-1)[: self.curve_dim]
|
|
|
if "curve_time" in s:
|
|
|
curve_time[i] = np.asarray(s["curve_time"], dtype=np.float32).reshape(-1)[: self.time_dim]
|
|
|
schedule_meta[i] = np.asarray(s["schedule_meta"], dtype=np.float32).reshape(-1)[: len(SCHEDULE_META_NAMES)]
|
|
|
family_name[i] = str(s.get("family_name", ""))
|
|
|
|
|
|
self.d_group_id[start:end] = gid
|
|
|
self.d_params[start:end] = params
|
|
|
self.d_schedule[start:end] = schedule
|
|
|
self.d_curve[start:end] = curve
|
|
|
self.d_curve_time[start:end] = curve_time
|
|
|
self.d_schedule_meta[start:end] = schedule_meta
|
|
|
self.d_family_name[start:end] = family_name
|
|
|
|
|
|
self._n = end
|
|
|
self.f.attrs["n_samples"] = int(self._n)
|
|
|
|
|
|
def flush(self):
|
|
|
self.f.flush()
|
|
|
|
|
|
def close(self):
|
|
|
self.flush()
|
|
|
self.f.close()
|
|
|
|
|
|
|
|
|
class ParallelDatasetGenerator:
|
|
|
"""大规模求解器生成数据集的并行调度层。"""
|
|
|
|
|
|
def __init__(self, cfg: Config, n_workers: int | None = None):
|
|
|
self.cfg = cfg
|
|
|
self.cfg.ensure_dirs()
|
|
|
self.n_workers = int(
|
|
|
n_workers or cfg.get("parallel", "n_workers", default=max(1, int(mp.cpu_count() * 0.75)))
|
|
|
)
|
|
|
self.output_dir = cfg.paths.samples_dir
|
|
|
|
|
|
def _generate_dataset(self) -> bool:
|
|
|
"""缺少 dataset.bin 时调用 C++ training.exe 自动生成。"""
|
|
|
try:
|
|
|
result = subprocess.run(
|
|
|
[str(self.cfg.training_exe)],
|
|
|
cwd=str(self.cfg.training_exe.parent),
|
|
|
timeout=60,
|
|
|
stdout=subprocess.DEVNULL,
|
|
|
stderr=subprocess.PIPE,
|
|
|
encoding="utf-8",
|
|
|
errors="ignore",
|
|
|
)
|
|
|
if result.returncode == 0 and self.cfg.dataset_bin.exists():
|
|
|
return True
|
|
|
|
|
|
fallback = self.cfg.training_exe.parent / "dataset.bin"
|
|
|
if result.returncode == 0 and fallback.exists():
|
|
|
self.cfg.dataset_bin.parent.mkdir(parents=True, exist_ok=True)
|
|
|
shutil.copy2(fallback, self.cfg.dataset_bin)
|
|
|
return True
|
|
|
|
|
|
return False
|
|
|
except Exception:
|
|
|
return False
|
|
|
|
|
|
def _cleanup_workers(self):
|
|
|
base_worker_dir = self.cfg.paths.temp_dir / "parallel_workers"
|
|
|
if base_worker_dir.exists():
|
|
|
shutil.rmtree(base_worker_dir, ignore_errors=True)
|
|
|
|
|
|
def generate(self, n_samples=None, method=None, random_seed=None, dataset_tag: str | None = None):
|
|
|
"""生成一份完整的正演代理模型训练 HDF5 数据集。"""
|
|
|
cfg = self.cfg
|
|
|
n_samples = int(n_samples or cfg.get("generation", "n_samples", default=50000))
|
|
|
method = method or cfg.raw["params"].get("sampling_method", "sobol")
|
|
|
random_seed = int(random_seed or cfg.get("generation", "random_seed", default=42))
|
|
|
|
|
|
if not cfg.dataset_bin.exists():
|
|
|
if bool(cfg.get("dataset_runtime", "auto_build_dataset_bin", default=True)):
|
|
|
if not self._generate_dataset():
|
|
|
raise RuntimeError("无法生成 dataset.bin")
|
|
|
else:
|
|
|
raise RuntimeError(f"dataset.bin 不存在: {cfg.dataset_bin}")
|
|
|
|
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
|
mode = str(cfg.raw["schedule"]["generation_mode"]).lower()
|
|
|
sec_mode = str(cfg.raw["schedule"]["section_policy"]["mode"]).lower()
|
|
|
tag = normalize_tag(dataset_tag)
|
|
|
if tag:
|
|
|
filepath = self.output_dir / f"dataset_{tag}_target{n_samples}_{timestamp}.h5"
|
|
|
else:
|
|
|
filepath = self.output_dir / f"dataset_{mode}_{sec_mode}_target{n_samples}_{timestamp}.h5"
|
|
|
|
|
|
curve_dim = cfg.curve_dim
|
|
|
Nu, Cu = cfg.schedule_grid_shape
|
|
|
sched_dim = Nu * Cu + cfg.sec_feat_dim
|
|
|
if bool(cfg.raw.get("schedule", {}).get("use_metadata_features_for_model", False)):
|
|
|
sched_dim += len(cfg.raw.get("schedule", {}).get("metadata_features_for_model", []) or [])
|
|
|
param_dim = 6
|
|
|
|
|
|
# HDF5 布局与正演模型输入输出一致：
|
|
|
# params/schedule 是输入，curve 是数值求解器生成的监督目标。
|
|
|
write_batch_size = int(cfg.get("streaming_hdf5", "write_batch_size", default=2000))
|
|
|
compression = cfg.get("streaming_hdf5", "compression", default=None)
|
|
|
chunk_rows = int(cfg.get("streaming_hdf5", "chunk_rows", default=2048))
|
|
|
|
|
|
app = HDF5Appender(
|
|
|
filepath=filepath,
|
|
|
param_dim=param_dim,
|
|
|
schedule_dim=sched_dim,
|
|
|
curve_dim=curve_dim,
|
|
|
compression=compression,
|
|
|
chunk_rows=chunk_rows,
|
|
|
)
|
|
|
|
|
|
successful = 0
|
|
|
total_requested = 0
|
|
|
fail_reasons = {}
|
|
|
fail_examples = {}
|
|
|
max_fail_examples_per_reason = int(cfg.get("generation", "max_fail_examples_per_reason", default=50))
|
|
|
max_rounds = int(cfg.get("generation", "max_rounds", default=999999))
|
|
|
max_total_requested = int(cfg.get("generation", "max_total_requested", default=int(n_samples * 12.0)))
|
|
|
|
|
|
checkpoint_every_n = int(cfg.get("parallel", "checkpoint_every_n", default=5000))
|
|
|
checkpoint_every_sec = float(cfg.get("parallel", "checkpoint_every_sec", default=180.0))
|
|
|
last_ckpt_n = 0
|
|
|
last_ckpt_ts = datetime.now().timestamp()
|
|
|
|
|
|
seed_base = int(random_seed)
|
|
|
round_id = 0
|
|
|
param_keys = cfg.raw["params"]["all_physical_param_names"]
|
|
|
expand_ma = 1.0
|
|
|
alpha = 0.2
|
|
|
expand_stats = []
|
|
|
buffer_samples = []
|
|
|
|
|
|
def flush_buffer(tag: str):
|
|
|
nonlocal buffer_samples, last_ckpt_n, last_ckpt_ts
|
|
|
if not buffer_samples:
|
|
|
return
|
|
|
|
|
|
# 写入进度元数据，便于长时间生成任务中断后检查当前状态。
|
|
|
app.append_samples(buffer_samples)
|
|
|
buffer_samples = []
|
|
|
app.flush()
|
|
|
|
|
|
meta = {
|
|
|
"is_partial": True,
|
|
|
"tag": tag,
|
|
|
"n_saved": int(app.n_samples),
|
|
|
"round_id": int(round_id),
|
|
|
"total_requested": int(total_requested),
|
|
|
"generated_at": datetime.now().isoformat(),
|
|
|
"generation_mode": mode,
|
|
|
"section_policy": sec_mode,
|
|
|
}
|
|
|
with open(filepath.with_suffix(".progress.json"), "w", encoding="utf-8") as f:
|
|
|
json.dump(meta, f, indent=2, ensure_ascii=False)
|
|
|
|
|
|
last_ckpt_n = int(app.n_samples)
|
|
|
last_ckpt_ts = datetime.now().timestamp()
|
|
|
|
|
|
try:
|
|
|
pbar = tqdm(total=n_samples, desc="valid samples", unit="sample")
|
|
|
|
|
|
while successful < n_samples:
|
|
|
round_id += 1
|
|
|
if round_id > max_rounds or total_requested >= max_total_requested:
|
|
|
break
|
|
|
|
|
|
need = n_samples - successful
|
|
|
request_params = int(np.ceil(need / max(expand_ma, 1.0) * 1.15))
|
|
|
request_params = max(1, min(request_params, 8000))
|
|
|
|
|
|
params_list = generate_params_dataset(
|
|
|
cfg,
|
|
|
request_params,
|
|
|
method=method,
|
|
|
random_seed=seed_base + round_id * 1000,
|
|
|
)
|
|
|
|
|
|
tasks = []
|
|
|
base_seed = seed_base + round_id * 1000
|
|
|
task_counter = 0
|
|
|
|
|
|
for pidx, p in enumerate(params_list):
|
|
|
pd_base = {k: float(getattr(p, k)) for k in param_keys}
|
|
|
schedule_id = int((round_id << 32) ^ (pidx << 1) ^ (base_seed & 0xFFFFFFFF))
|
|
|
rng_i = np.random.RandomState(_seed32(base_seed ^ (pidx + 12345)))
|
|
|
timeQ, q, sched_info = sample_schedule_by_mode(cfg, rng_i)
|
|
|
section_indices = _resolve_section_indices(cfg, timeQ, q, rng_i)
|
|
|
|
|
|
for sec in section_indices:
|
|
|
pd = dict(pd_base)
|
|
|
pd["schedule_id"] = schedule_id
|
|
|
pd["schedule"] = {
|
|
|
"sectionIndex": int(sec),
|
|
|
"timeQ": timeQ,
|
|
|
"q": q,
|
|
|
}
|
|
|
meta_vec, family_name = build_schedule_metadata(
|
|
|
cfg=cfg,
|
|
|
timeQ=timeQ,
|
|
|
q=q,
|
|
|
family_name=str(sched_info.get("family_name", "unknown")),
|
|
|
section_index=int(sec),
|
|
|
)
|
|
|
pd["schedule_meta"] = {
|
|
|
"meta_vec": meta_vec.tolist(),
|
|
|
"family_name": family_name,
|
|
|
}
|
|
|
tasks.append((task_counter, pd))
|
|
|
task_counter += 1
|
|
|
|
|
|
total_requested += len(tasks)
|
|
|
expand_now = len(tasks) / max(len(params_list), 1)
|
|
|
expand_ma = (1 - alpha) * expand_ma + alpha * expand_now
|
|
|
expand_stats.append(
|
|
|
{
|
|
|
"round": round_id,
|
|
|
"expand_now": float(expand_now),
|
|
|
"expand_ma": float(expand_ma),
|
|
|
}
|
|
|
)
|
|
|
|
|
|
reached = False
|
|
|
max_in_flight = int(cfg.get("parallel", "max_in_flight", default=self.n_workers * 4))
|
|
|
|
|
|
with ProcessPoolExecutor(
|
|
|
max_workers=self.n_workers,
|
|
|
initializer=_init_process_worker,
|
|
|
initargs=(str(cfg.path),),
|
|
|
) as executor:
|
|
|
task_iter = iter(tasks)
|
|
|
in_flight = deque()
|
|
|
|
|
|
for _ in range(min(max_in_flight, len(tasks))):
|
|
|
try:
|
|
|
task = next(task_iter)
|
|
|
in_flight.append(executor.submit(_worker_simulate_parallel, task))
|
|
|
except StopIteration:
|
|
|
break
|
|
|
|
|
|
while in_flight:
|
|
|
fut = in_flight.popleft()
|
|
|
try:
|
|
|
task_idx, sample, fail_reason, fail_ctx = fut.result()
|
|
|
except Exception:
|
|
|
fail_reasons["future_exception"] = fail_reasons.get("future_exception", 0) + 1
|
|
|
else:
|
|
|
if (not reached) and (sample is not None):
|
|
|
buffer_samples.append(sample)
|
|
|
successful += 1
|
|
|
pbar.update(1)
|
|
|
|
|
|
if len(buffer_samples) >= write_batch_size:
|
|
|
flush_buffer(f"batch_R{round_id}")
|
|
|
|
|
|
if successful >= n_samples:
|
|
|
reached = True
|
|
|
|
|
|
elif (not reached) and (sample is None) and fail_reason:
|
|
|
fail_reasons[fail_reason] = fail_reasons.get(fail_reason, 0) + 1
|
|
|
if len(fail_examples.get(fail_reason, [])) < max_fail_examples_per_reason:
|
|
|
fail_examples.setdefault(fail_reason, []).append(
|
|
|
{
|
|
|
"task_idx": int(task_idx),
|
|
|
"reason": str(fail_reason),
|
|
|
}
|
|
|
)
|
|
|
|
|
|
now = datetime.now().timestamp()
|
|
|
if ((app.n_samples + len(buffer_samples)) - last_ckpt_n >= checkpoint_every_n) or (
|
|
|
now - last_ckpt_ts >= checkpoint_every_sec
|
|
|
):
|
|
|
flush_buffer(f"ckpt_R{round_id}")
|
|
|
|
|
|
if not reached:
|
|
|
try:
|
|
|
task = next(task_iter)
|
|
|
in_flight.append(executor.submit(_worker_simulate_parallel, task))
|
|
|
except StopIteration:
|
|
|
pass
|
|
|
|
|
|
flush_buffer("final")
|
|
|
pbar.close()
|
|
|
|
|
|
with open(filepath.with_suffix(".fail_stats.json"), "w", encoding="utf-8") as f:
|
|
|
json.dump(
|
|
|
{
|
|
|
"fail_reasons": fail_reasons,
|
|
|
"fail_examples": fail_examples,
|
|
|
"expand_stats": expand_stats,
|
|
|
"total_requested": total_requested,
|
|
|
"n_success_target": int(n_samples),
|
|
|
"n_success_written": int(app.n_samples),
|
|
|
"generated_at": datetime.now().isoformat(),
|
|
|
},
|
|
|
f,
|
|
|
indent=2,
|
|
|
ensure_ascii=False,
|
|
|
)
|
|
|
|
|
|
self._cleanup_workers()
|
|
|
return str(filepath)
|
|
|
|
|
|
finally:
|
|
|
try:
|
|
|
if buffer_samples:
|
|
|
app.append_samples(buffer_samples)
|
|
|
app.close()
|
|
|
except Exception:
|
|
|
pass
|