Initial commit: photonAI MZM MLP baseline only

Made-with: Cursor

tests/test_smoke.py

@@ -0,0 +1,108 @@
+"""最小冒烟测试：模型 shape 与单 epoch 训练不报错。"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+from src.config import load_config
+from src.data import load_raw_txt
+from src.model import MLPRegressor
+from src.preprocess import prepare_training_data
+from src.trainer import fit
+
+
+def _write_synthetic_txt(path: Path, n: int = 64) -> None:
+    rng = np.random.default_rng(0)
+    x = rng.normal(size=(n, 8))
+    y = np.zeros((n, 3))
+    y[:, 0] = rng.normal(size=n)
+    y[:, 1] = rng.normal(size=n)
+    # 第 11 列 V_pi 落在默认物理门控 [0, 500] 内
+    y[:, 2] = rng.uniform(1.0, 400.0, size=n)
+    mat = np.hstack([x, y])
+    lines = [",".join(str(v) for v in row) for row in mat]
+    path.write_text("\n".join(lines), encoding="utf-8")
+
+
+def test_mlp_forward_shape() -> None:
+    m = MLPRegressor(8, [16, 16], 3, batchnorm=False, dropout=0.0, residual=False)
+    x = torch.randn(5, 8)
+    y = m(x)
+    assert y.shape == (5, 3)
+
+
+def test_one_epoch_training_pipeline(tmp_path: Path) -> None:
+    data_txt = tmp_path / "data.txt"
+    _write_synthetic_txt(data_txt, n=80)
+
+    cfg_dict = {
+        "data_path": str(data_txt),
+        "split_ratios": [0.7, 0.15, 0.15],
+        "random_seed": 1,
+        "remove_duplicate_rows": False,
+        "outlier_strategy": "none",
+        "outlier_apply_to": "targets",
+        "outlier_config": {
+            "iqr_k": 1.5,
+            "zscore_threshold": 4.0,
+            "quantile_lower": 0.001,
+            "quantile_upper": 0.999,
+        },
+        "remove_nonpositive_vpi": False,
+        "filter_v_pi_range": True,
+        "v_pi_min": 0.0,
+        "v_pi_max": 500.0,
+        "model": {
+            "input_dim": 8,
+            "hidden_dims": [32, 32],
+            "output_dim": 3,
+            "batchnorm": False,
+            "dropout": 0.0,
+            "residual": False,
+        },
+        "optimizer": {"name": "adamw", "lr": 0.01, "weight_decay": 0.0},
+        "scheduler": {
+            "type": "cosine",
+            "plateau_factor": 0.5,
+            "plateau_patience": 10,
+            "plateau_min_lr": 1e-6,
+        },
+        "training": {
+            "batch_size": 16,
+            "epochs": 1,
+            "early_stopping_patience": 1,
+            "num_workers": 0,
+        },
+        "loss": {"type": "huber", "huber_delta": 1.0, "target_weights": [1.0, 1.0, 1.0]},
+        "output_dir": str(tmp_path / "results"),
+    }
+    cfg_path = tmp_path / "cfg.yaml"
+    cfg_path.write_text(yaml.safe_dump(cfg_dict), encoding="utf-8")
+
+    cfg = load_config(cfg_path)
+    df = load_raw_txt(data_txt)
+
+    run_dir = tmp_path / "run0"
+    run_dir.mkdir()
+    (run_dir / "figures").mkdir()
+    (run_dir / "checkpoints").mkdir()
+
+    bundle = prepare_training_data(df, cfg, run_dir)
+    device = torch.device("cpu")
+    model = MLPRegressor(
+        input_dim=cfg.model.input_dim,
+        hidden_dims=cfg.model.hidden_dims,
+        output_dim=cfg.model.output_dim,
+        batchnorm=cfg.model.batchnorm,
+        dropout=cfg.model.dropout,
+        residual=cfg.model.residual,
+    ).to(device)
+    fit(model, cfg, bundle.train_loader, bundle.val_loader, run_dir, device)
+    assert (run_dir / "checkpoints" / "best.pt").is_file()
+    meta = json.loads((run_dir / "cleaning_meta.json").read_text(encoding="utf-8"))
+    assert meta["n_train"] > 0