118 lines
3.7 KiB
Python
118 lines
3.7 KiB
Python
"""最小冒烟测试：MoE + PINN 流程可运行。"""
|
||
|
||
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 create_model_from_config
|
||
from src.preprocess import inverse_transform_targets, prepare_training_data
|
||
from src.trainer import compute_pinn_loss, fit
|
||
|
||
|
||
def _write_synthetic_txt(path: Path, n: int = 96) -> None:
|
||
rng = np.random.default_rng(0)
|
||
x = rng.normal(size=(n, 8))
|
||
y = np.zeros((n, 3))
|
||
length = np.abs(x[:, 7]) + 0.5
|
||
y[:, 0] = 2.0 - 0.2 * length + rng.normal(scale=0.05, size=n)
|
||
y[:, 1] = 0.3 + 0.1 * length + rng.normal(scale=0.03, size=n)
|
||
y[:, 2] = 20.0 / length + rng.normal(scale=0.2, size=n)
|
||
mat = np.hstack([x, y])
|
||
path.write_text("\n".join(",".join(str(v) for v in row) for row in mat), encoding="utf-8")
|
||
|
||
|
||
def _make_cfg(tmp_path: Path, data_txt: Path, epochs: int = 1) -> Path:
|
||
cfg_dict = {
|
||
"data_path": str(data_txt),
|
||
"data": {
|
||
"test_size": 0.1,
|
||
"random_state": 123,
|
||
"filter_v_pi_max": 500.0,
|
||
},
|
||
"model": {
|
||
"input_dim": 8,
|
||
"output_dim": 3,
|
||
"hidden_dims": [16, 16],
|
||
"n_experts": 4,
|
||
"gating_hidden": 4,
|
||
"dropout_rate": 0.0,
|
||
"use_bn": False,
|
||
"activation": "relu",
|
||
},
|
||
"optimizer": {
|
||
"lr": 0.001,
|
||
"weight_decay": 0.01,
|
||
"betas": [0.9, 0.999],
|
||
},
|
||
"training": {
|
||
"batch_size": 16,
|
||
"epochs": epochs,
|
||
"num_workers": 0,
|
||
},
|
||
"physics": {
|
||
"lambda_bw_mon": 0.1,
|
||
"lambda_IL_mon": 0.1,
|
||
"lambda_vpiL": 0.05,
|
||
"lambda_smooth": 0.01,
|
||
},
|
||
"output_dir": str(tmp_path / "results"),
|
||
}
|
||
cfg_path = tmp_path / "cfg.yaml"
|
||
cfg_path.write_text(yaml.safe_dump(cfg_dict), encoding="utf-8")
|
||
return cfg_path
|
||
|
||
|
||
def test_moe_forward_shape(tmp_path: Path) -> None:
|
||
data_txt = tmp_path / "data.txt"
|
||
_write_synthetic_txt(data_txt, n=32)
|
||
cfg = load_config(_make_cfg(tmp_path, data_txt))
|
||
model = create_model_from_config(cfg)
|
||
x = torch.randn(5, 8)
|
||
y = model(x)
|
||
assert y.shape == (5, 3)
|
||
|
||
|
||
def test_pinn_loss_backpropagates(tmp_path: Path) -> None:
|
||
data_txt = tmp_path / "data.txt"
|
||
_write_synthetic_txt(data_txt, n=40)
|
||
cfg = load_config(_make_cfg(tmp_path, data_txt))
|
||
model = create_model_from_config(cfg)
|
||
x = torch.randn(8, 8)
|
||
y = torch.randn(8, 3)
|
||
loss, data_loss, terms = compute_pinn_loss(model, x, y, torch.nn.MSELoss(), cfg)
|
||
loss.backward()
|
||
assert float(loss.detach()) >= float(data_loss.detach())
|
||
assert "IL_mon" in terms
|
||
assert any(p.grad is not None for p in model.parameters())
|
||
|
||
|
||
def test_one_epoch_training_pipeline(tmp_path: Path) -> None:
|
||
data_txt = tmp_path / "data.txt"
|
||
_write_synthetic_txt(data_txt, n=80)
|
||
cfg = load_config(_make_cfg(tmp_path, data_txt, epochs=1))
|
||
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 = create_model_from_config(cfg).to(device)
|
||
history = fit(model, cfg, bundle, run_dir, device)
|
||
assert (run_dir / "checkpoints" / "last.pt").is_file()
|
||
assert len(history.train_loss) == 1
|
||
meta = json.loads((run_dir / "cleaning_meta.json").read_text(encoding="utf-8"))
|
||
assert meta["n_train"] > 0
|
||
assert meta["n_test"] > 0
|
||
restored = inverse_transform_targets(bundle.y_test[:3], bundle.y_scalers)
|
||
assert restored.shape == (3, 3)
|