64 lines
2.1 KiB
Python
64 lines
2.1 KiB
Python
import os
|
|
from pathlib import Path
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
import yaml
|
|
|
|
|
|
def load_config(path: str) -> dict:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return yaml.safe_load(f)
|
|
|
|
|
|
def save_checkpoint(model, optimizer: torch.optim.Optimizer, iteration: int, path: str):
|
|
os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
|
|
torch.save(
|
|
{
|
|
"iteration": iteration,
|
|
"model_state_dict": model.state_dict(),
|
|
"optimizer_state_dict": optimizer.state_dict(),
|
|
},
|
|
path,
|
|
)
|
|
print(f"[Checkpoint] saved → {path}")
|
|
|
|
|
|
def load_checkpoint(model, path: str, device=None) -> int:
|
|
ckpt = torch.load(path, map_location=device or "cpu")
|
|
model.load_state_dict(ckpt["model_state_dict"], strict=False)
|
|
if "optimizer_state_dict" in ckpt and hasattr(model, "optimizer"):
|
|
try:
|
|
model.optimizer.load_state_dict(ckpt["optimizer_state_dict"])
|
|
except Exception:
|
|
pass
|
|
it = ckpt.get("iteration", 0)
|
|
print(f"[Checkpoint] loaded ← {path} (iter {it})")
|
|
return it
|
|
|
|
|
|
def setup_helmholtz_config(config: dict, k_test=None, center=None, radius=None, eps_test=None) -> float:
|
|
"""Lock scatterer/helmholtz config for test/viz. Returns wave number k."""
|
|
hc = config.setdefault("environment", {}).setdefault("mesh_refinement", {}).setdefault("fem", {}).setdefault("helmholtz", {})
|
|
sc = hc.setdefault("scatterer", {})
|
|
sc["mode"] = "fixed"
|
|
if center is not None:
|
|
sc["cx"], sc["cy"] = center[0], center[1]
|
|
if radius is not None:
|
|
sc["radius"] = radius
|
|
if eps_test is not None:
|
|
sc["eps_r"] = eps_test
|
|
if k_test is not None:
|
|
hc["wave_number_mode"] = "fixed"
|
|
hc["wave_number"] = k_test
|
|
return hc.get("wave_number", 6.0)
|
|
|
|
|
|
def parse_center(center_str: Optional[str]) -> Optional[Tuple[float, float]]:
|
|
if center_str is None:
|
|
return None
|
|
parts = center_str.split(",")
|
|
if len(parts) != 2:
|
|
raise ValueError(f"Invalid --center format (expected 'cx,cy'): {center_str}")
|
|
return (float(parts[0].strip()), float(parts[1].strip()))
|