examples/schedulers/slurm_dtlz2_from_design_json.py

0001 #!/usr/bin/env python
0002 """Compute DTLZ2 objectives from a design-point JSON file and write JSON output."""
0003
0004 from __future__ import annotations
0005
0006 import argparse
0007 import json
0008 import math
0009 import sys
0010 import time
0011 from pathlib import Path
0012 from typing import Any, Dict
0013
0014
0015 def parse_args() -> argparse.Namespace:
0016     parser = argparse.ArgumentParser(description=__doc__)
0017     parser.add_argument("--design", type=Path, required=True)
0018     parser.add_argument("--output", type=Path, required=True)
0019     parser.add_argument("--label", type=str, default="slurm-dtlz2")
0020     parser.add_argument("--sleep-sec", type=float, default=0.0)
0021     parser.add_argument("--repeat", type=int, default=2)
0022     return parser.parse_args()
0023
0024
0025 def load_design(path: Path) -> Dict[str, Any]:
0026     payload = json.loads(path.read_text(encoding="utf-8"))
0027     required = [
0028         "DTLZ2_variables.x1",
0029         "DTLZ2_variables.x2",
0030         "DTLZ2_variables.x3",
0031         "DTLZ2_variables.x4",
0032         "DTLZ2_variables.x5",
0033     ]
0034     missing = [key for key in required if key not in payload]
0035     if missing:
0036         raise ValueError(f"Design file {path} is missing required keys: {missing}")
0037     return payload
0038
0039
0040 def compute_objectives(parameters: Dict[str, Any]) -> Dict[str, float]:
0041     x1 = float(parameters["DTLZ2_variables.x1"])
0042     tail = [
0043         float(parameters["DTLZ2_variables.x2"]),
0044         float(parameters["DTLZ2_variables.x3"]),
0045         float(parameters["DTLZ2_variables.x4"]),
0046         float(parameters["DTLZ2_variables.x5"]),
0047     ]
0048     g = sum((value - 0.5) ** 2 for value in tail)
0049     factor = 1.0 + g
0050     f1 = factor * math.cos(x1 * math.pi / 2.0)
0051     f2 = factor * math.sin(x1 * math.pi / 2.0)
0052     return {"f1": float(f1), "f2": float(f2)}
0053
0054
0055 def main() -> int:
0056     args = parse_args()
0057     design = load_design(args.design)
0058
0059     print(f"[{args.label}] design file: {args.design}", flush=True)
0060     for step in range(args.repeat):
0061         print(f"[{args.label}] progress {step + 1}/{args.repeat}", flush=True)
0062         print(f"[{args.label}] stderr heartbeat {step + 1}", file=sys.stderr, flush=True)
0063         if args.sleep_sec > 0:
0064             time.sleep(args.sleep_sec)
0065
0066     payload = compute_objectives(design)
0067     args.output.parent.mkdir(parents=True, exist_ok=True)
0068     args.output.write_text(json.dumps(payload, indent=2), encoding="utf-8")
0069     print(json.dumps(payload, indent=2), flush=True)
0070     return 0
0071
0072
0073 if __name__ == "__main__":
0074     raise SystemExit(main())