|
|
|
|
|
"""批量生成正演代理模型的原始数值试井数据集。
|
|
|
|
|
|
|
|
|
|
|
|
脚本读取数据生成配置，调用并行数据集生成器批量采样地层/井筒参数和流量制度，
|
|
|
|
|
|
运行底层数值求解器并把有效曲线写入 HDF5。它是训练前最上游的数据生产入口。
|
|
|
|
|
|
"""
|
|
|
|
|
|
|
|
|
|
|
|
from __future__ import annotations
|
|
|
|
|
|
import argparse
|
|
|
|
|
|
import multiprocessing as mp
|
|
|
|
|
|
import sys
|
|
|
|
|
|
from pathlib import Path
|
|
|
|
|
|
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
|
|
|
|
sys.path.append(str(ROOT))
|
|
|
|
|
|
|
|
|
|
|
|
from src.common.config import Config
|
|
|
|
|
|
from src.common.experiment_paths import config_for_stage
|
|
|
|
|
|
from src.data.dataset_generation import ParallelDatasetGenerator
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def main():
|
|
|
|
|
|
"""按配置阶段启动并行数值试井样本生成，输出原始 HDF5 数据集路径。"""
|
|
|
|
|
|
parser = argparse.ArgumentParser()
|
|
|
|
|
|
parser.add_argument("--config", default=None)
|
|
|
|
|
|
parser.add_argument(
|
|
|
|
|
|
"--stage",
|
|
|
|
|
|
choices=["fixed_case", "case_neighborhood", "family_random", "family_random_hard", "family_random_v2_q"],
|
|
|
|
|
|
default=None,
|
|
|
|
|
|
)
|
|
|
|
|
|
parser.add_argument("--n-samples", type=int, default=None)
|
|
|
|
|
|
parser.add_argument("--n-workers", type=int, default=None)
|
|
|
|
|
|
parser.add_argument("--seed", type=int, default=None)
|
|
|
|
|
|
parser.add_argument("--method", type=str, default=None)
|
|
|
|
|
|
parser.add_argument("--dataset-tag", type=str, default=None, help="Optional tag injected into output dataset filename")
|
|
|
|
|
|
args = parser.parse_args()
|
|
|
|
|
|
|
|
|
|
|
|
config_path = args.config
|
|
|
|
|
|
if config_path is None:
|
|
|
|
|
|
config_path = str(config_for_stage(args.stage) or Path("configs/data_gen.yaml"))
|
|
|
|
|
|
|
|
|
|
|
|
# stage 用来选择预设配置；命令行参数继续覆盖样本数、并行数和随机种子。
|
|
|
|
|
|
cfg = Config(config_path)
|
|
|
|
|
|
cfg.ensure_dirs()
|
|
|
|
|
|
path = ParallelDatasetGenerator(cfg=cfg, n_workers=args.n_workers).generate(
|
|
|
|
|
|
n_samples=args.n_samples, method=args.method, random_seed=args.seed, dataset_tag=args.dataset_tag
|
|
|
|
|
|
)
|
|
|
|
|
|
print(path)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
|
|
mp.freeze_support()
|
|
|
|
|
|
try:
|
|
|
|
|
|
mp.set_start_method("spawn", force=True)
|
|
|
|
|
|
except Exception:
|
|
|
|
|
|
pass
|
|
|
|
|
|
main()
|
