diff --git a/.gitignore b/.gitignore index ea0f6eb84..8e4cc2ff9 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,5 @@ logs/ serialized_dataset/ exporter.log installation_DOE_supercomputers/HydraGNN-Installation-Frontier*/ +*.doc +*.docx diff --git a/examples/multidataset_hpo_sc26/job-omnistat-deephyper.sh b/examples/multidataset_hpo_sc26/job-omnistat-deephyper.sh index 432634235..4ac1df942 100755 --- a/examples/multidataset_hpo_sc26/job-omnistat-deephyper.sh +++ b/examples/multidataset_hpo_sc26/job-omnistat-deephyper.sh @@ -40,7 +40,7 @@ function setup_bb() fi } -HYDRAGNN_ROOT=/lustre/orion/world-shared/lrn070/jyc/frontier/HydraGNN +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN # Load conda environment module reset diff --git a/examples/multidataset_hpo_sc26/job-omnistat_train_single_model.sh b/examples/multidataset_hpo_sc26/job-omnistat_train_single_model.sh index 339578f78..e3a25a545 100755 --- a/examples/multidataset_hpo_sc26/job-omnistat_train_single_model.sh +++ b/examples/multidataset_hpo_sc26/job-omnistat_train_single_model.sh @@ -13,7 +13,7 @@ function cmd() { time $@ } -HYDRAGNN_ROOT=/lustre/orion/world-shared/lrn070/jyc/frontier/HydraGNN +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN # Load conda environment module reset diff --git a/examples/multidataset_hpo_sc26/job-sc26-single-model-training-frontier.sh b/examples/multidataset_hpo_sc26/job-sc26-single-model-training-frontier.sh deleted file mode 100755 index 9c6b23b10..000000000 --- a/examples/multidataset_hpo_sc26/job-sc26-single-model-training-frontier.sh +++ /dev/null @@ -1,163 +0,0 @@ -#!/bin/bash -#SBATCH -A LRN070 -#SBATCH -J HydraGNN -#SBATCH -o job-%j.out -#SBATCH -e job-%j.out -#SBATCH -t 02:00:00 -#SBATCH -N 16 -#SBATCH -p batch -#SBATCH -q debug -#SBATCH -C nvme - -export all_proxy=socks://proxy.ccs.ornl.gov:3128/ -export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ -export http_proxy=http://proxy.ccs.ornl.gov:3128/ -export https_proxy=http://proxy.ccs.ornl.gov:3128/ -export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' - -function cmd() { - echo "$@" - time $@ -} - -function setup_bb() -{ - # Move a copy of the env to the NVMe on each node - if [ -d /mnt/bb/${USER} ]; then - srun -N${SLURM_JOB_NUM_NODES} --ntasks-per-node 1 -l -u df -h /mnt/bb/${USER} | grep -v Filesystem - SRCDIR=/lustre/orion/world-shared/lrn070/jyc/frontier/HydraGNN - for FILENAME in HydraGNN-Installation-Frontier.tar; do - echo "Checking ${FILENAME}" - if [ ! -f /mnt/bb/${USER}/${FILENAME} ]; then - echo "Copying ${FILENAME} to each local NVME" - time sbcast -pfv ${SRCDIR}/${FILENAME} /mnt/bb/${USER}/${FILENAME} - if [ ! "$?" == "0" ]; then - # CHECK EXIT CODE. When SBCAST fails, it may leave partial files on the compute nodes, and if you continue to launch srun, - # your application may pick up partially complete shared library files, which would give you confusing errors. - echo "SBCAST failed!" - break - fi - fi - echo "Untar ${FILENAME}" - time srun -N${SLURM_JOB_NUM_NODES} --ntasks-per-node 1 -l -u tar -xf /mnt/bb/${USER}/${FILENAME} -C /mnt/bb/${USER}/ - done - echo "NVME is ready to use" - srun -N${SLURM_JOB_NUM_NODES} --ntasks-per-node 1 -l -u df -h /mnt/bb/${USER} | grep -v Filesystem - fi -} - -HYDRAGNN_ROOT=/lustre/orion/mat746/proj-shared/mlupopa/HydraGNN - -# Load conda environemnt -source /lustre/orion/mat746/proj-shared/mlupopa/module-to-load-frontier-rocm640.sh - -source activate /lustre/orion/mat746/proj-shared/mlupopa/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier-ROCm6.4/hydragnn_venv - - -# setup_bb -# if [ -d /mnt/bb/${USER}/HydraGNN-Installation-Frontier ]; then -# export PYTHONPATH=/mnt/bb/${USER}/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH -# fi - -#export python path to HydragNN -export PYTHONPATH=$PWD:$PYTHONPATH - -echo "===== Module List =====" -module list - -echo "===== Check =====" -which python -python -c "import adios2; print(adios2.__version__, adios2.__file__)" -python -c "import torch; print(torch.__version__, torch.__file__)" - -echo "===== LD_LIBRARY_PATH =====" -echo $LD_LIBRARY_PATH | tr ':' '\n' - -export MPICH_ENV_DISPLAY=0 -export MPICH_VERSION_DISPLAY=0 -export MIOPEN_DISABLE_CACHE=1 -export MIOPEN_USER_DB_PATH=/tmp/$JOBID -mkdir -p $MIOPEN_USER_DB_PATH - -export PYTHONNOUSERSITE=1 - -export OMP_NUM_THREADS=7 -export HYDRAGNN_NUM_WORKERS=0 -export HYDRAGNN_USE_VARIABLE_GRAPH_SIZE=1 -export HYDRAGNN_AGGR_BACKEND=mpi -export HYDRAGNN_VALTEST=1 - -## Checking -env | grep ROCM -env | grep ^MI -env | grep ^MPICH -env | grep ^HYDRA - -export HYDRAGNN_TRACE_LEVEL=0 -export HYDRAGNN_MAX_NUM_BATCH=1000 -export TASK_PARALLEL=0 -export HYDRAGNN_TASK_PARALLEL_PROPORTIONAL_SPLIT=0 -export BATCH_SIZE=40 -export NUM_EPOCH=50 - -export HYDRAGNN_DDSTORE_METHOD=1 -export HYDRAGNN_CUSTOM_DATALOADER=1 -export HYDRAGNN_NUM_WORKERS=1 - -# Dataset ordering matches gfm_deephyper_multi_all_mpnn.py multi_model_list -export datadir0=Alexandria -export datadir1=ANI1x -export datadir2=MPTrj -export datadir3=OC2020 -export datadir4=OC2022 -export datadir5=OC25 -export datadir6=ODAC23 -export datadir7=OMat24 -export datadir8=OMol25 -export datadir9=OMol25-neutral -export datadir10=OMol25-non-neutral -export datadir11=OPoly2026 -export datadir12=Nabla2DFT -export datadir13=QCML -export datadir14=QM7X -export datadir15=transition1x - - -# (A) Setup omnistat sampling environment -ml use /sw/frontier/amdsw/modulefiles/ -ml omnistat-wrapper -export OMNISTAT_CONFIG=$HYDRAGNN_ROOT/omnistat.hydragnn-external-fp64.config - -# (B) Enable data collectors and polling (1 sec interval) -${OMNISTAT_WRAPPER} usermode --start --interval 15 - -## Verification mode for this script: FSDP2 + SHARD_GRAD_OP, task-parallel optional -export HYDRAGNN_USE_FSDP=1 -export HYDRAGNN_FSDP_VERSION=2 -export HYDRAGNN_FSDP_STRATEGY=SHARD_GRAD_OP -TASK_PARALLEL_ARG="" -if [ "$TASK_PARALLEL" = "1" ]; then - TASK_PARALLEL_ARG="--task_parallel" -fi - -## Getting strange omnistat counts with the following line. Otherswise, FSDP runs fine. - -MULTI_MODEL_LIST=$datadir0 -# MULTI_MODEL_LIST=$datadir0,$datadir1,$datadir2,$datadir3,$datadir4,$datadir5,$datadir6,$datadir7,$datadir8,$datadir9,$datadir10,$datadir11,$datadir12,$datadir13,$datadir14,$datadir15 - -DATASET=datadir$K -cmd srun -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 --gpus-per-task=1 --gpu-bind=closest -l --kill-on-bad-exit=1 \ -python -u $HYDRAGNN_ROOT/examples/multidataset_hpo_sc26/gfm_mlip_all_mpnn.py \ - --log=multidataset_hpo-$SLURM_JOB_ID-NN$SLURM_JOB_NUM_NODES-FSDP$HYDRAGNN_USE_FSDP-V$HYDRAGNN_FSDP_VERSION-TP$TASK_PARALLEL --everyone \ - --inputfile=gfm_mlip.json --num_samples=$((BATCH_SIZE*HYDRAGNN_MAX_NUM_BATCH*NUM_EPOCH)) \ - --multi --ddstore --multi_model_list=$MULTI_MODEL_LIST --batch_size=$BATCH_SIZE --num_epoch=$NUM_EPOCH \ - $TASK_PARALLEL_ARG \ - --precision=fp64 \ - --mpnn_type=EGNN \ - --num_conv_layers=2 \ - --hidden_dim=1000 \ - --num_headlayers=2 \ - --dim_headlayers=300 - -# (C) End of job: stop data collection -${OMNISTAT_WRAPPER} usermode --stop diff --git a/examples/opf/__init__.py b/examples/opf/__init__.py new file mode 100644 index 000000000..d60edd97b --- /dev/null +++ b/examples/opf/__init__.py @@ -0,0 +1,26 @@ +from dataclasses import dataclass + +from download_and_uncompress_data import ( + _discover_cases, + _ensure_opf_downloaded, + _parse_num_groups, + _resolve_num_groups, +) + + +@dataclass(frozen=True) +class _DataOps: + discover_cases: callable + ensure_opf_downloaded: callable + parse_num_groups: callable + resolve_num_groups: callable + + +data_ops = _DataOps( + discover_cases=_discover_cases, + ensure_opf_downloaded=_ensure_opf_downloaded, + parse_num_groups=_parse_num_groups, + resolve_num_groups=_resolve_num_groups, +) + +__all__ = ["data_ops"] diff --git a/examples/opf/check_domain_loss_on_gt.py b/examples/opf/check_domain_loss_on_gt.py new file mode 100644 index 000000000..03c92efcf --- /dev/null +++ b/examples/opf/check_domain_loss_on_gt.py @@ -0,0 +1,178 @@ +"""Check whether domain loss penalties are zero on ground-truth OPF solutions. + +Loads N raw JSON samples, substitutes bus.y (ground truth Va, Vm) as the +"prediction", and evaluates each penalty term. A correctly-posed feasibility +penalty should return exactly zero (or near-zero up to float precision) on +any strictly feasible OPF solution. + +Usage: + python3 check_domain_loss_on_gt.py [--n_samples N] [--json_dir PATH] +""" + +import argparse +import json +import os +import sys + +import torch +import torch.nn.functional as F + +# --------------------------------------------------------------------------- +# Parse args +# --------------------------------------------------------------------------- +parser = argparse.ArgumentParser() +parser.add_argument("--n_samples", type=int, default=20) +parser.add_argument( + "--json_dir", + default=( + "dataset/dataset_release_1/pglib_opf_case10000_goc/raw/" + "gridopt-dataset-tmp/dataset_release_1/pglib_opf_case10000_goc/group_1" + ), +) +args = parser.parse_args() + +# --------------------------------------------------------------------------- +# Minimal HeteroData builder (same logic as _raw_json_to_heterodata) +# --------------------------------------------------------------------------- +try: + import torch_geometric.datasets.opf as tg_opf + from torch_geometric.data import HeteroData +except ImportError: + sys.exit("torch_geometric not available") + + +def load_sample(filepath): + with open(filepath) as f: + obj = json.load(f) + grid = obj["grid"] + solution = obj["solution"] + + data = HeteroData() + data["bus"].x = torch.tensor(grid["nodes"]["bus"], dtype=torch.float) + data["bus"].y = torch.tensor(solution["nodes"]["bus"], dtype=torch.float) + + def _ei(obj, rel): + try: + return tg_opf.extract_edge_index(obj, rel) + except Exception: + return None + + for rel in ("ac_line", "transformer"): + ei = _ei(obj, rel) + if ei is None: + continue + ea = grid["edges"].get(rel, {}).get("features") + if ea is None: + continue + data["bus", rel, "bus"].edge_index = ei + data["bus", rel, "bus"].edge_attr = torch.tensor(ea, dtype=torch.float) + + return data + + +# --------------------------------------------------------------------------- +# Penalty evaluation on ground-truth bus.y +# --------------------------------------------------------------------------- +VMIN_IDX, VMAX_IDX = 2, 3 # bus.x columns for v_min, v_max +VM_IDX = 1 # bus.y column for Vm +VA_IDX = 0 # bus.y column for Va + +# ac_line edge_attr indices +AC_X_IDX, AC_RATE_A_IDX = 5, 6 +AC_THETA_MIN_IDX, AC_THETA_MAX_IDX = 0, 1 + +# transformer edge_attr indices +TR_X_IDX, TR_RATE_A_IDX = 3, 4 +TR_THETA_MIN_IDX, TR_THETA_MAX_IDX = 0, 1 + + +def evaluate(data): + bus_true = data["bus"].y # shape [N_bus, 2] (Va, Vm) + bus_x = data["bus"].x # shape [N_bus, 4] + + results = {} + + # 1. Voltage bound: vmin <= Vm <= vmax + Vm = bus_true[:, VM_IDX] + vmin = bus_x[:, VMIN_IDX] + vmax = bus_x[:, VMAX_IDX] + vbound = torch.mean(F.relu(vmin - Vm).pow(2) + F.relu(Vm - vmax).pow(2)) + results["voltage_bound"] = vbound.item() + + # 2. Angle difference limit: theta_min <= Va_i - Va_j <= theta_max + Va = bus_true[:, VA_IDX] + for rel, tmin_i, tmax_i, tag in [ + (("bus", "ac_line", "bus"), AC_THETA_MIN_IDX, AC_THETA_MAX_IDX, "ac_angle_diff"), + (("bus", "transformer", "bus"), TR_THETA_MIN_IDX, TR_THETA_MAX_IDX, "tr_angle_diff"), + ]: + if rel not in data.edge_types: + continue + ea = getattr(data[rel], "edge_attr", None) + ei = getattr(data[rel], "edge_index", None) + if ea is None or ei is None or ea.shape[1] <= max(tmin_i, tmax_i): + continue + theta_min = ea[:, tmin_i] + theta_max = ea[:, tmax_i] + src, dst = ei + delta = Va[src] - Va[dst] + p = torch.mean(F.relu(delta - theta_max).pow(2) + F.relu(theta_min - delta).pow(2)) + results[tag] = p.item() + + # 3. DC thermal limit: |P_ij| = |(Va_i - Va_j) / x_ij| <= rate_a + for rel, x_i, ra_i, tag in [ + (("bus", "ac_line", "bus"), AC_X_IDX, AC_RATE_A_IDX, "ac_line_flow"), + (("bus", "transformer", "bus"), TR_X_IDX, TR_RATE_A_IDX, "tr_line_flow"), + ]: + if rel not in data.edge_types: + continue + ea = getattr(data[rel], "edge_attr", None) + ei = getattr(data[rel], "edge_index", None) + if ea is None or ei is None or ea.shape[1] <= max(x_i, ra_i): + continue + x_ij = ea[:, x_i].clamp(min=1e-6) + rate_a = ea[:, ra_i].clamp(min=0.0) + src, dst = ei + P_ij = (Va[src] - Va[dst]) / x_ij + p = torch.mean(F.relu(P_ij.abs() - rate_a).pow(2)) + results[tag] = p.item() + + return results + + +# --------------------------------------------------------------------------- +# Run over N samples +# --------------------------------------------------------------------------- +json_dir = os.path.join(os.path.dirname(__file__), args.json_dir) +files = sorted( + [os.path.join(json_dir, f) for f in os.listdir(json_dir) if f.endswith(".json")] +)[: args.n_samples] + +if not files: + sys.exit(f"No JSON files found in {json_dir}") + +print(f"Evaluating {len(files)} samples from {json_dir}\n") + +accum = {} +n_ok = 0 +for path in files: + try: + data = load_sample(path) + r = evaluate(data) + for k, v in r.items(): + accum.setdefault(k, []).append(v) + n_ok += 1 + except Exception as e: + print(f" WARNING: {os.path.basename(path)}: {e}") + +print(f"{'Term':<25} {'mean':>12} {'max':>12} {'min':>12} {'#non-zero':>10}") +print("-" * 75) +eps = 1e-9 +for term in sorted(accum): + vals = accum[term] + mean_v = sum(vals) / len(vals) + max_v = max(vals) + min_v = min(vals) + nonzero = sum(1 for v in vals if abs(v) > eps) + print(f"{term:<25} {mean_v:>12.6e} {max_v:>12.6e} {min_v:>12.6e} {nonzero:>10}/{len(vals)}") + +print(f"\nProcessed {n_ok}/{len(files)} samples successfully.") diff --git a/examples/opf/compare_heat_runs.py b/examples/opf/compare_heat_runs.py new file mode 100644 index 000000000..b4efada22 --- /dev/null +++ b/examples/opf/compare_heat_runs.py @@ -0,0 +1,121 @@ +#!/usr/bin/env python3 +"""Compare two HydraGNN training logs (baseline vs physics-informed). + +Usage: + python compare_heat_runs.py [--json] +""" +import json +import re +import sys +from pathlib import Path + +PATTERN = re.compile( + r"Epoch:\s*(\d+),\s*Train Loss:\s*([0-9eE+\-.]+),\s*Val Loss:\s*([0-9eE+\-.]+),\s*Test Loss:\s*([0-9eE+\-.]+)" +) + + +def extract_stats(log_path: Path): + """Return first, best, last epoch rows plus total epoch count, or None.""" + rows = [] + with log_path.open() as f: + for line in f: + m = PATTERN.search(line) + if not m: + continue + epoch, train_loss, val_loss, test_loss = m.groups() + rows.append( + { + "epoch": int(epoch), + "train_loss": float(train_loss), + "val_loss": float(val_loss), + "test_loss": float(test_loss), + } + ) + if not rows: + return None + best = min(rows, key=lambda r: r["val_loss"]) + return { + "first": rows[0], + "last": rows[-1], + "best": best, + "num_epochs": len(rows), + } + + +def _fmt(stats: dict) -> str: + b = stats["best"] + return ( + f" epochs logged : {stats['num_epochs']}\n" + f" first epoch : {stats['first']['epoch']} " + f"train={stats['first']['train_loss']:.6f} " + f"val={stats['first']['val_loss']:.6f} " + f"test={stats['first']['test_loss']:.6f}\n" + f" last epoch : {stats['last']['epoch']} " + f"train={stats['last']['train_loss']:.6f} " + f"val={stats['last']['val_loss']:.6f} " + f"test={stats['last']['test_loss']:.6f}\n" + f" best val epoch: {b['epoch']} " + f"train={b['train_loss']:.6f} " + f"val={b['val_loss']:.6f} " + f"test={b['test_loss']:.6f}" + ) + + +def main(argv): + use_json = "--json" in argv + paths = [a for a in argv[1:] if not a.startswith("--")] + if len(paths) != 2: + raise SystemExit( + "Usage: compare_heat_runs.py [--json]" + ) + + baseline_path = Path(paths[0]) + physics_path = Path(paths[1]) + + for p in (baseline_path, physics_path): + if not p.exists(): + raise SystemExit(f"Log file not found: {p}") + + baseline = extract_stats(baseline_path) + physics = extract_stats(physics_path) + + if baseline is None: + raise SystemExit(f"No epoch metrics found in baseline log: {baseline_path}") + if physics is None: + raise SystemExit(f"No epoch metrics found in physics log: {physics_path}") + + delta_val = physics["best"]["val_loss"] - baseline["best"]["val_loss"] + delta_test = physics["best"]["test_loss"] - baseline["best"]["test_loss"] + + if use_json: + print( + json.dumps( + { + "baseline": baseline, + "physics": physics, + "delta_best_val_loss": delta_val, + "delta_best_test_loss": delta_test, + }, + indent=2, + ) + ) + return + + print(f"Baseline ({baseline_path}):") + print(_fmt(baseline)) + print() + print(f"Physics-informed ({physics_path}):") + print(_fmt(physics)) + print() + sign_val = "+" if delta_val >= 0 else "" + sign_test = "+" if delta_test >= 0 else "" + print(f"Delta best val_loss (physics - baseline): {sign_val}{delta_val:.6f}") + print(f"Delta best test_loss (physics - baseline): {sign_test}{delta_test:.6f}") + if delta_val < 0: + print("=> Physics-informed loss IMPROVED best validation loss.") + else: + print("=> Physics-informed loss did NOT improve best validation loss.") + + +if __name__ == "__main__": + main(sys.argv) diff --git a/examples/opf/download_and_uncompress_data.py b/examples/opf/download_and_uncompress_data.py new file mode 100644 index 000000000..b53cf04a0 --- /dev/null +++ b/examples/opf/download_and_uncompress_data.py @@ -0,0 +1,446 @@ +import logging +import os +import shutil +import socket +import urllib.error +import urllib.request +from contextlib import contextmanager + +import torch_geometric.datasets.opf as tg_opf + + +def _opf_release_name(topological_perturbations: bool) -> str: + return ( + "dataset_release_1_nminusone" + if topological_perturbations + else "dataset_release_1" + ) + + +def _opf_raw_dir(root: str, case_name: str, topological_perturbations: bool) -> str: + return os.path.join( + root, _opf_release_name(topological_perturbations), case_name, "raw" + ) + + +def _opf_release_dir(root: str, topological_perturbations: bool) -> str: + return os.path.join(root, _opf_release_name(topological_perturbations)) + + +def _opf_tmp_dir(root: str, case_name: str, topological_perturbations: bool) -> str: + return os.path.join( + _opf_raw_dir(root, case_name, topological_perturbations), "gridopt-dataset-tmp" + ) + + +def _opf_group_dir( + root: str, + case_name: str, + group_idx: int, + topological_perturbations: bool, +) -> str: + return os.path.join( + _opf_tmp_dir(root, case_name, topological_perturbations), + _opf_release_name(topological_perturbations), + case_name, + f"group_{group_idx}", + ) + + +def _opf_remote_url( + case_name: str, group_idx: int, topological_perturbations: bool +) -> str: + release = _opf_release_name(topological_perturbations) + return ( + "https://storage.googleapis.com/gridopt-dataset/" + f"{release}/{case_name}_{group_idx}.tar.gz" + ) + + +def _is_address_family_not_supported(exc: Exception) -> bool: + err = exc + if isinstance(exc, urllib.error.URLError): + err = getattr(exc, "reason", exc) + return getattr(err, "errno", None) == 97 + + +@contextmanager +def _force_ipv4_getaddrinfo(): + original_getaddrinfo = socket.getaddrinfo + + def _ipv4_getaddrinfo(host, port, family=0, type=0, proto=0, flags=0): + if family in (0, socket.AF_UNSPEC): + family = socket.AF_INET + return original_getaddrinfo(host, port, family, type, proto, flags) + + socket.getaddrinfo = _ipv4_getaddrinfo + try: + yield + finally: + socket.getaddrinfo = original_getaddrinfo + + +def _urlopen_with_ipv4_fallback(request_or_url, timeout: int): + try: + return urllib.request.urlopen(request_or_url, timeout=timeout) + except urllib.error.URLError as exc: + if not _is_address_family_not_supported(exc): + raise + logging.warning( + "URL open hit Errno 97 (address family not supported); retrying with IPv4" + ) + + with _force_ipv4_getaddrinfo(): + return urllib.request.urlopen(request_or_url, timeout=timeout) + + +def _url_exists(url: str, timeout: int = 10) -> bool: + req = urllib.request.Request(url, method="HEAD") + try: + with _urlopen_with_ipv4_fallback(req, timeout=timeout) as resp: + return 200 <= resp.status < 300 + except urllib.error.HTTPError as exc: + if exc.code == 404: + return False + raise + except urllib.error.URLError as exc: + raise RuntimeError(f"Unable to reach {url}: {exc}") from exc + + +def _download_file(url: str, dst_path: str, timeout: int = 60) -> None: + os.makedirs(os.path.dirname(dst_path), exist_ok=True) + tmp_path = f"{dst_path}.part.{os.getpid()}" + try: + with _urlopen_with_ipv4_fallback(url, timeout=timeout) as response: + with open(tmp_path, "wb") as out_file: + shutil.copyfileobj(response, out_file) + os.replace(tmp_path, dst_path) + except Exception: + if os.path.exists(tmp_path): + try: + os.remove(tmp_path) + except OSError: + pass + raise + + +def _parallel_download_and_extract_opf( + root, + case_name, + num_groups, + topological_perturbations, + rank, + comm, +): + raw_dir = _opf_raw_dir(root, case_name, topological_perturbations) + os.makedirs(raw_dir, exist_ok=True) + + world_size = comm.Get_size() + active_workers = min(world_size, num_groups) + if rank >= active_workers: + return + + assigned_groups = list(range(rank, num_groups, active_workers)) + logging.info( + "OPF parallel fetch/extract: rank=%d assigned_groups=%d", + rank, + len(assigned_groups), + ) + + skipped_download = 0 + skipped_extract = 0 + downloaded = 0 + extracted = 0 + + for group_idx in assigned_groups: + archive_name = f"{case_name}_{group_idx}.tar.gz" + archive_path = os.path.join(raw_dir, archive_name) + group_dir = _opf_group_dir( + root, case_name, group_idx, topological_perturbations + ) + + if not os.path.isfile(archive_path): + archive_url = _opf_remote_url( + case_name, group_idx, topological_perturbations + ) + logging.info( + "OPF parallel fetch: rank=%d case=%s group=%d action=download", + rank, + case_name, + group_idx, + ) + _download_file(archive_url, archive_path) + downloaded += 1 + else: + skipped_download += 1 + + if os.path.isdir(group_dir): + skipped_extract += 1 + continue + + logging.info( + "OPF parallel extract: rank=%d case=%s group=%d action=extract", + rank, + case_name, + group_idx, + ) + tg_opf.extract_tar(archive_path, raw_dir) + extracted += 1 + + if skipped_download or skipped_extract: + logging.info( + "OPF parallel fetch/extract summary: rank=%d case=%s groups=%d downloaded=%d extracted=%d skip_download=%d skip_extract=%d", + rank, + case_name, + len(assigned_groups), + downloaded, + extracted, + skipped_download, + skipped_extract, + ) + + +def _probe_remote_num_groups( + case_name: str, + topological_perturbations: bool, + start_idx: int, + max_groups: int, +) -> int: + count = max(0, int(start_idx)) + for idx in range(count, max_groups): + url = _opf_remote_url(case_name, idx, topological_perturbations) + if _url_exists(url): + count = idx + 1 + continue + break + return count + + +def _find_empty_json(root: str): + empty = [] + for dirpath, _, filenames in os.walk(root): + for name in filenames: + if not name.endswith(".json"): + continue + path = os.path.join(dirpath, name) + try: + if os.path.getsize(path) == 0: + empty.append(path) + except OSError: + empty.append(path) + return empty + + +def _discover_cases(root: str, topological_perturbations: bool): + release_dir = _opf_release_dir(root, topological_perturbations) + if not os.path.isdir(release_dir): + return [] + return sorted( + name + for name in os.listdir(release_dir) + if os.path.isdir(os.path.join(release_dir, name)) + ) + + +def _discover_num_groups(root: str, case_name: str, topological_perturbations: bool): + raw_dir = _opf_raw_dir(root, case_name, topological_perturbations) + if not os.path.isdir(raw_dir): + return 0 + groups = [] + for name in os.listdir(raw_dir): + if not name.startswith(f"{case_name}_") or not name.endswith(".tar.gz"): + continue + try: + idx = int(name[len(case_name) + 1 : -len(".tar.gz")]) + groups.append(idx) + except ValueError: + continue + return max(groups) + 1 if groups else 0 + + +def _parse_num_groups(num_groups_arg: str) -> int | None: + if isinstance(num_groups_arg, int): + return num_groups_arg + if str(num_groups_arg).lower() == "all": + return None + return int(num_groups_arg) + + +def _resolve_num_groups( + requested_num_groups, + datadir, + case_name, + topological_perturbations, + num_groups_max, + probe_remote, + rank, + comm, +): + if requested_num_groups is not None: + return requested_num_groups + local_groups = _discover_num_groups(datadir, case_name, topological_perturbations) + if not probe_remote: + if local_groups == 0: + if num_groups_max <= 0: + raise RuntimeError(f"No groups found for case '{case_name}'.") + return num_groups_max + return local_groups + if num_groups_max <= 0: + raise RuntimeError(f"No groups found for case '{case_name}'.") + if local_groups >= num_groups_max: + return local_groups + if rank == 0: + logging.info( + "Probing remote groups for %s starting at %d (cap %d)", + case_name, + local_groups, + num_groups_max, + ) + try: + resolved = _probe_remote_num_groups( + case_name, + topological_perturbations, + local_groups, + num_groups_max, + ) + except RuntimeError as exc: + if local_groups > 0: + logging.warning( + "Remote group probe failed for %s; using local groups=%d. " + "To disable remote probe, pass --no_num_groups_probe. Error: %s", + case_name, + local_groups, + exc, + ) + resolved = local_groups + else: + raise + if resolved == 0: + raise RuntimeError(f"No groups found for case '{case_name}'.") + else: + resolved = None + return comm.bcast(resolved, root=0) + + +def _reextract_opf_if_needed(root, case_name, num_groups, topological_perturbations): + raw_dir = _opf_raw_dir(root, case_name, topological_perturbations) + tmp_dir = _opf_tmp_dir(root, case_name, topological_perturbations) + raw_files = [f"{case_name}_{i}.tar.gz" for i in range(num_groups)] + + if not os.path.isdir(raw_dir): + logging.info( + "OPF extract check: skip case=%s reason=raw_dir_missing path=%s", + case_name, + raw_dir, + ) + return + + missing = [ + name for name in raw_files if not os.path.isfile(os.path.join(raw_dir, name)) + ] + if missing: + logging.info( + "OPF extract check: skip case=%s reason=missing_archives count=%d", + case_name, + len(missing), + ) + return + + if not os.path.isdir(tmp_dir): + logging.info( + "OPF extract check: skip case=%s reason=tmp_dir_missing (already processed)", + case_name, + ) + return + + empty = _find_empty_json(tmp_dir) + if not empty: + logging.info( + "OPF extract check: skip case=%s reason=tmp_dir_healthy path=%s", + case_name, + tmp_dir, + ) + return + logging.warning( + "OPF extract check: reextract case=%s reason=empty_json count=%d", + case_name, + len(empty), + ) + shutil.rmtree(tmp_dir, ignore_errors=True) + + logging.info( + "OPF extract check: extracting case=%s groups=%d", + case_name, + num_groups, + ) + for name in raw_files: + tg_opf.extract_tar(os.path.join(raw_dir, name), raw_dir) + + +def _ensure_missing_group_dirs( + root, + case_name, + num_groups, + topological_perturbations, +): + raw_dir = _opf_raw_dir(root, case_name, topological_perturbations) + missing_groups = [] + for group_idx in range(num_groups): + group_dir = _opf_group_dir( + root, case_name, group_idx, topological_perturbations + ) + if not os.path.isdir(group_dir): + missing_groups.append(group_idx) + + if not missing_groups: + return + + logging.warning( + "OPF extract check: missing group dirs case=%s count=%d; re-extracting archives", + case_name, + len(missing_groups), + ) + for group_idx in missing_groups: + archive_path = os.path.join(raw_dir, f"{case_name}_{group_idx}.tar.gz") + if os.path.isfile(archive_path): + tg_opf.extract_tar(archive_path, raw_dir) + + +def _ensure_opf_downloaded( + root, + case_name, + num_groups, + topological_perturbations, + rank, + comm, +): + _parallel_download_and_extract_opf( + root, + case_name, + num_groups, + topological_perturbations, + rank, + comm, + ) + comm.Barrier() + + if rank == 0: + _reextract_opf_if_needed( + root, + case_name, + num_groups, + topological_perturbations, + ) + _ensure_missing_group_dirs( + root, + case_name, + num_groups, + topological_perturbations, + ) + logging.info( + "OPF dataset ready: case=%s groups=%s topological_perturbations=%s", + case_name, + num_groups, + topological_perturbations, + ) + comm.Barrier() diff --git a/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_full.json b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_full.json new file mode 100644 index 000000000..0ddbb2778 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_full.json @@ -0,0 +1,111 @@ +{ + "_ft_strategy": "FT1_feasibility_classification", + "_ft_description": "Feasibility classification: binary graph-level prediction of whether an OPF instance is feasible (1) or infeasible (0). Infeasible samples are synthesised by scaling load features by an overload factor so that total demand exceeds generation capacity. Preprocessing: run generate_infeasible_samples.py to create the mixed dataset before training.", + "_ft_overload_factor": 6.0, + "ft_data_modelname": "FT1_feasibility_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "graph": [ + { + "type": "branch-0", + "architecture": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ] + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 1 + ], + "output_type": [ + "graph" + ], + "num_nodes": null, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "output_names": [ + "feasibility" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 1 + ], + "type": [ + "graph" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "binary_cross_entropy", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_head_only.json b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_head_only.json new file mode 100644 index 000000000..cb2da25a6 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_head_only.json @@ -0,0 +1,111 @@ +{ + "_ft_strategy": "FT1_feasibility_classification", + "_ft_description": "Feasibility classification: binary graph-level prediction of whether an OPF instance is feasible (1) or infeasible (0). Infeasible samples are synthesised by scaling load features by an overload factor so that total demand exceeds generation capacity. Preprocessing: run generate_infeasible_samples.py to create the mixed dataset before training.", + "_ft_overload_factor": 6.0, + "ft_data_modelname": "FT1_feasibility_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "graph": [ + { + "type": "branch-0", + "architecture": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ] + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 1 + ], + "output_type": [ + "graph" + ], + "num_nodes": null, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "output_names": [ + "feasibility" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 1 + ], + "type": [ + "graph" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "binary_cross_entropy", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_partial.json b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_partial.json new file mode 100644 index 000000000..8856b00c5 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroHEAT_partial.json @@ -0,0 +1,111 @@ +{ + "_ft_strategy": "FT1_feasibility_classification", + "_ft_description": "Feasibility classification: binary graph-level prediction of whether an OPF instance is feasible (1) or infeasible (0). Infeasible samples are synthesised by scaling load features by an overload factor so that total demand exceeds generation capacity. Preprocessing: run generate_infeasible_samples.py to create the mixed dataset before training.", + "_ft_overload_factor": 6.0, + "ft_data_modelname": "FT1_feasibility_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "graph": [ + { + "type": "branch-0", + "architecture": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ] + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 1 + ], + "output_type": [ + "graph" + ], + "num_nodes": null, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "output_names": [ + "feasibility" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 1 + ], + "type": [ + "graph" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "binary_cross_entropy", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_full.json b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_full.json new file mode 100644 index 000000000..b80e0c682 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_full.json @@ -0,0 +1,111 @@ +{ + "_ft_strategy": "FT1_feasibility_classification", + "_ft_description": "Feasibility classification: binary graph-level prediction of whether an OPF instance is feasible (1) or infeasible (0). Infeasible samples are synthesised by scaling load features by an overload factor so that total demand exceeds generation capacity. Preprocessing: run generate_infeasible_samples.py to create the mixed dataset before training.", + "_ft_overload_factor": 6.0, + "ft_data_modelname": "FT1_feasibility_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "graph": [ + { + "type": "branch-0", + "architecture": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ] + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 1 + ], + "output_type": [ + "graph" + ], + "num_nodes": null, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "output_names": [ + "feasibility" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 1 + ], + "type": [ + "graph" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "binary_cross_entropy", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_head_only.json b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_head_only.json new file mode 100644 index 000000000..6da43e9f3 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_head_only.json @@ -0,0 +1,111 @@ +{ + "_ft_strategy": "FT1_feasibility_classification", + "_ft_description": "Feasibility classification: binary graph-level prediction of whether an OPF instance is feasible (1) or infeasible (0). Infeasible samples are synthesised by scaling load features by an overload factor so that total demand exceeds generation capacity. Preprocessing: run generate_infeasible_samples.py to create the mixed dataset before training.", + "_ft_overload_factor": 6.0, + "ft_data_modelname": "FT1_feasibility_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "graph": [ + { + "type": "branch-0", + "architecture": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ] + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 1 + ], + "output_type": [ + "graph" + ], + "num_nodes": null, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "output_names": [ + "feasibility" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 1 + ], + "type": [ + "graph" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "binary_cross_entropy", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_partial.json b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_partial.json new file mode 100644 index 000000000..b58162342 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/config_HeteroSAGE_partial.json @@ -0,0 +1,111 @@ +{ + "_ft_strategy": "FT1_feasibility_classification", + "_ft_description": "Feasibility classification: binary graph-level prediction of whether an OPF instance is feasible (1) or infeasible (0). Infeasible samples are synthesised by scaling load features by an overload factor so that total demand exceeds generation capacity. Preprocessing: run generate_infeasible_samples.py to create the mixed dataset before training.", + "_ft_overload_factor": 6.0, + "ft_data_modelname": "FT1_feasibility_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "graph": [ + { + "type": "branch-0", + "architecture": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ] + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 1 + ], + "output_type": [ + "graph" + ], + "num_nodes": null, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "output_names": [ + "feasibility" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 1 + ], + "type": [ + "graph" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "binary_cross_entropy", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-HeteroHEAT.sh b/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-HeteroHEAT.sh new file mode 100644 index 000000000..180a7f4a7 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-HeteroHEAT.sh @@ -0,0 +1,152 @@ +#!/bin/bash +# ============================================================================= +# Frontier Slurm job — FT1 Feasibility Classification (HeteroHEAT) +# +# FT1 pipeline differs from FT2-FT4 in two ways: +# Phase 1 — generate the mixed feasible/infeasible HDF5 dataset with +# generate_infeasible_samples.py (runs once, shared by both archs) +# Phase 2 — fine-tune with train_opf_ft1_classify.py (graph-level BCE) +# +# NOTE: If HeteroSAGE Phase 1 has already run (FT1_feasibility_data.h5 +# already exists), set PHASES=train to skip regeneration. +# +# Usage: +# sbatch job-frontier-FT1-HeteroHEAT.sh +# +# Override regime or phases: +# sbatch --export=ALL,FT_REGIME=partial,PHASES=train job-frontier-FT1-HeteroHEAT.sh +# ============================================================================= +#SBATCH -A LRN078 +#SBATCH -J OPF-FT1-HEAT +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/finetune-FT1-HEAT-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/finetune-FT1-HEAT-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- +FT_ARCH=HeteroHEAT +FT_REGIME=${FT_REGIME:-full} +PRETRAINED_MODEL=${PRETRAINED_MODEL:-HeteroHEAT_best} +PHASES=${PHASES:-preonly,train} +# Set NO_PRETRAINED=1 to train from random initialisation (baseline) +NO_PRETRAINED=${NO_PRETRAINED:-0} + +FT1_DATASET=FT1_feasibility_data +SRC_DATASET=${SRC_DATASET:-FT3_contingency_data} +OVERLOAD_FACTOR=${OVERLOAD_FACTOR:-6.0} +MAX_SAMPLES=${MAX_SAMPLES:-5000} + +N_TRAIN=${SLURM_JOB_NUM_NODES} + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +# ----------------------------------------------------------------------------- +# Derived paths +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +DATA_ROOT=$FT_DIR/../dataset +_n_tag=${MAX_TRAIN_SAMPLES:+_n${MAX_TRAIN_SAMPLES}} +_scratch_tag=${NO_PRETRAINED:+_scratch} +LOG_NAME="FT1_feasibility_${FT_ARCH}_${FT_REGIME}${_scratch_tag}${_n_tag}" + +SRC_DIR=$DATA_ROOT/${SRC_DATASET}.h5 +OUT_DIR=$DATA_ROOT/${FT1_DATASET}.h5 + +echo "============================================================" +echo " FT1 Feasibility Classification on Frontier" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Regime : $FT_REGIME" +echo " Phases : $PHASES" +echo " Source dataset : $SRC_DATASET.h5" +echo " FT1 dataset : $FT1_DATASET.h5" +echo " Overload factor : $OVERLOAD_FACTOR" +echo " Max samples : $MAX_SAMPLES (per class)" +echo " Log name : $LOG_NAME" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +# ============================================================================= +# Phase 1: Generate mixed feasible/infeasible dataset +# ============================================================================= +if [[ "$PHASES" == *"preonly"* ]]; then + echo "" + echo "--- Phase 1: Generating FT1 feasibility dataset ---" + + if [ -d "$OUT_DIR" ]; then + echo " Dataset already exists at $OUT_DIR — skipping generation." + echo " Delete $OUT_DIR to regenerate." + else + srun -N1 -n1 -c7 \ + python -u generate_infeasible_samples.py \ + --src_dir $SRC_DIR \ + --out_dir $OUT_DIR \ + --overload_factor $OVERLOAD_FACTOR \ + --max_samples $MAX_SAMPLES + echo "--- Phase 1 complete: $(ls $OUT_DIR) ---" + fi +fi + +# ============================================================================= +# Phase 2: Fine-tune for binary classification +# ============================================================================= +if [[ "$PHASES" == *"train"* ]]; then + echo "" + echo "--- Phase 2: FT1 classification fine-tuning ($FT_REGIME) with $FT_ARCH ---" + srun -N$N_TRAIN -n$((N_TRAIN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_ft1_classify.py \ + --inputfile FT1_feasibility_classification/config_${FT_ARCH}_${FT_REGIME}.json \ + --modelname $LOG_NAME \ + --data_root $DATA_ROOT \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $FT_REGIME \ + ${NO_PRETRAINED:+--no_pretrained} \ + ${MAX_TRAIN_SAMPLES:+--max_train_samples $MAX_TRAIN_SAMPLES} + echo "--- Phase 2 complete ---" +fi + +echo "" +echo "FT1 job finished. Logs: $HYDRAGNN_ROOT/logs/$LOG_NAME/" diff --git a/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-HeteroSAGE.sh b/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-HeteroSAGE.sh new file mode 100644 index 000000000..51f843eb5 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-HeteroSAGE.sh @@ -0,0 +1,160 @@ +#!/bin/bash +# ============================================================================= +# Frontier Slurm job — FT1 Feasibility Classification (HeteroSAGE) +# +# FT1 pipeline differs from FT2-FT4 in two ways: +# Phase 1 — generate the mixed feasible/infeasible HDF5 dataset with +# generate_infeasible_samples.py (runs once, shared by both archs) +# Phase 2 — fine-tune with train_opf_ft1_classify.py (graph-level BCE) +# +# Usage: +# sbatch job-frontier-FT1-HeteroSAGE.sh +# +# Override regime or phases from CLI: +# sbatch --export=ALL,FT_REGIME=partial,PHASES=train job-frontier-FT1-HeteroSAGE.sh +# ============================================================================= +#SBATCH -A LRN078 +#SBATCH -J OPF-FT1-SAGE +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/finetune-FT1-SAGE-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/finetune-FT1-SAGE-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- +FT_ARCH=HeteroSAGE +FT_REGIME=${FT_REGIME:-full} # full | partial | head_only +PRETRAINED_MODEL=${PRETRAINED_MODEL:-HeteroSAGE_best} +PHASES=${PHASES:-preonly,train} # preonly,train | preonly | train +# Set NO_PRETRAINED=1 to train from random initialisation (baseline) +NO_PRETRAINED=${NO_PRETRAINED:-0} + +# FT1 dataset name (shared across both archs) +FT1_DATASET=FT1_feasibility_data + +# Source dataset for infeasible-sample generation. +# Defaults to FT3 contingency data (case118, N-1) — a diverse set of +# feasible samples. Change to any other preprocessed OPF HDF5 directory. +SRC_DATASET=${SRC_DATASET:-FT3_contingency_data} + +# Factor by which load features (Pd, Qd) are multiplied to make instances +# infeasible. Values >=5 reliably exceed generation capacity for pglib cases. +OVERLOAD_FACTOR=${OVERLOAD_FACTOR:-6.0} + +# Maximum feasible samples to use (same number of infeasible samples generated) +MAX_SAMPLES=${MAX_SAMPLES:-5000} + +N_TRAIN=${SLURM_JOB_NUM_NODES} + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +# ----------------------------------------------------------------------------- +# Derived paths +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +DATA_ROOT=$FT_DIR/../dataset +_n_tag=${MAX_TRAIN_SAMPLES:+_n${MAX_TRAIN_SAMPLES}} +_scratch_tag=${NO_PRETRAINED:+_scratch} +LOG_NAME="FT1_feasibility_${FT_ARCH}_${FT_REGIME}${_scratch_tag}${_n_tag}" + +SRC_DIR=$DATA_ROOT/${SRC_DATASET}.h5 +OUT_DIR=$DATA_ROOT/${FT1_DATASET}.h5 + +echo "============================================================" +echo " FT1 Feasibility Classification on Frontier" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Regime : $FT_REGIME" +echo " Phases : $PHASES" +echo " Source dataset : $SRC_DATASET.h5" +echo " FT1 dataset : $FT1_DATASET.h5" +echo " Overload factor : $OVERLOAD_FACTOR" +echo " Max samples : $MAX_SAMPLES (per class)" +echo " Log name : $LOG_NAME" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +# ============================================================================= +# Phase 1: Generate mixed feasible/infeasible dataset +# ============================================================================= +if [[ "$PHASES" == *"preonly"* ]]; then + echo "" + echo "--- Phase 1: Generating FT1 feasibility dataset ---" + + if [ -d "$OUT_DIR" ]; then + echo " Dataset already exists at $OUT_DIR — skipping generation." + echo " Delete $OUT_DIR to regenerate." + else + # Single-rank generation (no MPI parallelism needed) + srun -N1 -n1 -c7 \ + python -u generate_infeasible_samples.py \ + --src_dir $SRC_DIR \ + --out_dir $OUT_DIR \ + --overload_factor $OVERLOAD_FACTOR \ + --max_samples $MAX_SAMPLES + echo "--- Phase 1 complete: $(ls $OUT_DIR) ---" + fi +fi + +# ============================================================================= +# Phase 2: Fine-tune for binary classification +# ============================================================================= +if [[ "$PHASES" == *"train"* ]]; then + echo "" + echo "--- Phase 2: FT1 classification fine-tuning ($FT_REGIME) with $FT_ARCH ---" + srun -N$N_TRAIN -n$((N_TRAIN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_ft1_classify.py \ + --inputfile FT1_feasibility_classification/config_${FT_ARCH}_${FT_REGIME}.json \ + --modelname $LOG_NAME \ + --data_root $DATA_ROOT \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $FT_REGIME \ + ${NO_PRETRAINED:+--no_pretrained} \ + ${MAX_TRAIN_SAMPLES:+--max_train_samples $MAX_TRAIN_SAMPLES} + echo "--- Phase 2 complete ---" +fi + +echo "" +echo "FT1 job finished. Logs: $HYDRAGNN_ROOT/logs/$LOG_NAME/" diff --git a/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-multi.sh b/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-multi.sh new file mode 100644 index 000000000..16118f484 --- /dev/null +++ b/examples/opf/finetune/FT1_feasibility_classification/job-frontier-FT1-multi.sh @@ -0,0 +1,133 @@ +#!/bin/bash +# ============================================================================= +# Frontier Slurm job — FT1 Feasibility Classification (multi-method) +# +# Runs all 4 training methods in parallel within a single allocation: +# - full fine-tuning (all layers) +# - partial fine-tuning (last conv + head) +# - head-only fine-tuning (linear probe) +# - from-scratch baseline (random init) +# +# Each method gets N_PER_RUN nodes; total allocation = 4 × N_PER_RUN. +# +# Required env vars (set via sbatch --export or environment): +# FT_ARCH HeteroSAGE | HeteroHEAT +# +# Optional env vars: +# MAX_TRAIN_SAMPLES limit training samples (default: use all) +# N_PER_RUN nodes per method (default: 8) +# PRETRAINED_MODEL override pretrained model name +# +# Usage (from examples/opf/finetune/): +# sbatch --export=ALL,FT_ARCH=HeteroSAGE,MAX_TRAIN_SAMPLES=5000 \ +# FT1_feasibility_classification/job-frontier-FT1-multi.sh +# ============================================================================= +#SBATCH -A LRN070 +#SBATCH -J OPF-FT1-MULTI +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT1-multi-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT1-multi-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -N 32 + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- +FT_ARCH=${FT_ARCH:-HeteroSAGE} +MAX_TRAIN_SAMPLES=${MAX_TRAIN_SAMPLES:-} +N_PER_RUN=${N_PER_RUN:-8} +PRETRAINED_MODEL=${PRETRAINED_MODEL:-${FT_ARCH}_best} + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:${PYTHONPATH:-} + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +# ----------------------------------------------------------------------------- +# Derived +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +DATA_ROOT=$FT_DIR/../dataset +_n_tag=${MAX_TRAIN_SAMPLES:+_n${MAX_TRAIN_SAMPLES}} + +echo "============================================================" +echo " FT1 Feasibility Classification — multi-method job" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Methods : full / partial / head_only / scratch (parallel)" +echo " Nodes per run : $N_PER_RUN (total: $((N_PER_RUN * 4)))" +echo " Max train samples: ${MAX_TRAIN_SAMPLES:-all}" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +# Helper: launch one srun in the background, redirect output to its own log +_launch() { + local REGIME="$1" + local SCRATCH_FLAG="${2:-}" + local _scratch_tag="" + [[ -n "$SCRATCH_FLAG" ]] && _scratch_tag="_scratch" + local LOG_NAME="FT1_feasibility_${FT_ARCH}_${REGIME}${_scratch_tag}${_n_tag}" + local LOGFILE="$HYDRAGNN_ROOT/${LOG_NAME}-${SLURM_JOB_ID}.out" + + echo " Launching $LOG_NAME → $LOGFILE" + + srun --exact -N${N_PER_RUN} -n$((N_PER_RUN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_ft1_classify.py \ + --inputfile FT1_feasibility_classification/config_${FT_ARCH}_${REGIME}.json \ + --modelname "$LOG_NAME" \ + --data_root $DATA_ROOT \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $REGIME \ + $SCRATCH_FLAG \ + ${MAX_TRAIN_SAMPLES:+--max_train_samples $MAX_TRAIN_SAMPLES} \ + > "$LOGFILE" 2>&1 & +} + +# Launch all 4 methods in parallel +_launch full +_launch partial +_launch head_only +_launch full "--no_pretrained" # scratch baseline + +echo "" +echo "All 4 methods launched. Waiting for completion..." +wait +echo "" +echo "All methods finished. Job ID: $SLURM_JOB_ID" diff --git a/examples/opf/finetune/FT1_topology/config_HeteroHEAT_full.json b/examples/opf/finetune/FT1_topology/config_HeteroHEAT_full.json new file mode 100644 index 000000000..72ee66a8b --- /dev/null +++ b/examples/opf/finetune/FT1_topology/config_HeteroHEAT_full.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT1_topology", + "_ft_description": "Topology-specific fine-tuning: pretrain on 10-case corpus, fine-tune on held-out topology pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT1_topology_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_topology/config_HeteroHEAT_head_only.json b/examples/opf/finetune/FT1_topology/config_HeteroHEAT_head_only.json new file mode 100644 index 000000000..ca91e1bba --- /dev/null +++ b/examples/opf/finetune/FT1_topology/config_HeteroHEAT_head_only.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT1_topology", + "_ft_description": "Topology-specific fine-tuning: pretrain on 10-case corpus, fine-tune on held-out topology pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT1_topology_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_topology/config_HeteroHEAT_partial.json b/examples/opf/finetune/FT1_topology/config_HeteroHEAT_partial.json new file mode 100644 index 000000000..cab048b2b --- /dev/null +++ b/examples/opf/finetune/FT1_topology/config_HeteroHEAT_partial.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT1_topology", + "_ft_description": "Topology-specific fine-tuning: pretrain on 10-case corpus, fine-tune on held-out topology pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT1_topology_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_topology/config_HeteroSAGE_full.json b/examples/opf/finetune/FT1_topology/config_HeteroSAGE_full.json new file mode 100644 index 000000000..49c5bdfec --- /dev/null +++ b/examples/opf/finetune/FT1_topology/config_HeteroSAGE_full.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT1_topology", + "_ft_description": "Topology-specific fine-tuning: pretrain on 10-case corpus, fine-tune on held-out topology pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT1_topology_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_topology/config_HeteroSAGE_head_only.json b/examples/opf/finetune/FT1_topology/config_HeteroSAGE_head_only.json new file mode 100644 index 000000000..5ccd2437f --- /dev/null +++ b/examples/opf/finetune/FT1_topology/config_HeteroSAGE_head_only.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT1_topology", + "_ft_description": "Topology-specific fine-tuning: pretrain on 10-case corpus, fine-tune on held-out topology pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT1_topology_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_topology/config_HeteroSAGE_partial.json b/examples/opf/finetune/FT1_topology/config_HeteroSAGE_partial.json new file mode 100644 index 000000000..3681aaff9 --- /dev/null +++ b/examples/opf/finetune/FT1_topology/config_HeteroSAGE_partial.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT1_topology", + "_ft_description": "Topology-specific fine-tuning: pretrain on 10-case corpus, fine-tune on held-out topology pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT1_topology_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT1_topology/job-frontier-HeteroHEAT.sh b/examples/opf/finetune/FT1_topology/job-frontier-HeteroHEAT.sh new file mode 100755 index 000000000..6fdfa3275 --- /dev/null +++ b/examples/opf/finetune/FT1_topology/job-frontier-HeteroHEAT.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Convenience launcher for FT1_topology with HeteroHEAT. +# Usage: sbatch FT1_topology/job-frontier-HeteroHEAT.sh +# Override regime via SBATCH --export, e.g.: +# sbatch --export=ALL,FT_REGIME=partial FT1_topology/job-frontier-HeteroHEAT.sh +#SBATCH -A LRN078 +#SBATCH -J OPF-FT1-HeteroHEAT +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT1_topology-HeteroHEAT-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT1_topology-HeteroHEAT-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +export FT_STRATEGY=FT1_topology +export FT_ARCH=HeteroHEAT +export FT_REGIME=${FT_REGIME:-full} +export PRETRAINED_MODEL=HeteroHEAT_best + +bash $(dirname $0)/../job-frontier-finetune.sh diff --git a/examples/opf/finetune/FT1_topology/job-frontier-HeteroSAGE.sh b/examples/opf/finetune/FT1_topology/job-frontier-HeteroSAGE.sh new file mode 100755 index 000000000..e40443e1f --- /dev/null +++ b/examples/opf/finetune/FT1_topology/job-frontier-HeteroSAGE.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Convenience launcher for FT1_topology with HeteroSAGE. +# Usage: sbatch FT1_topology/job-frontier-HeteroSAGE.sh +# Override regime via SBATCH --export, e.g.: +# sbatch --export=ALL,FT_REGIME=partial FT1_topology/job-frontier-HeteroSAGE.sh +#SBATCH -A LRN078 +#SBATCH -J OPF-FT1-HeteroSAGE +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT1_topology-HeteroSAGE-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT1_topology-HeteroSAGE-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +export FT_STRATEGY=FT1_topology +export FT_ARCH=HeteroSAGE +export FT_REGIME=${FT_REGIME:-full} +export PRETRAINED_MODEL=HeteroSAGE_best + +bash $(dirname $0)/../job-frontier-finetune.sh diff --git a/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_full.json b/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_full.json new file mode 100644 index 000000000..3d31e0588 --- /dev/null +++ b/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_full.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT2_operating_condition", + "_ft_description": "Operating-condition fine-tuning: fine-tune on a new distribution of load/generation profiles for pglib_opf_case14_ieee with a limited labeled budget (5000 samples).", + "_ft_case_name": "pglib_opf_case14_ieee", + "_ft_num_groups": "2", + "_ft_max_samples": 5000, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT2_operating_condition_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_head_only.json b/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_head_only.json new file mode 100644 index 000000000..f0525eb4c --- /dev/null +++ b/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_head_only.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT2_operating_condition", + "_ft_description": "Operating-condition fine-tuning: fine-tune on a new distribution of load/generation profiles for pglib_opf_case14_ieee with a limited labeled budget (5000 samples).", + "_ft_case_name": "pglib_opf_case14_ieee", + "_ft_num_groups": "2", + "_ft_max_samples": 5000, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT2_operating_condition_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_partial.json b/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_partial.json new file mode 100644 index 000000000..30dc19796 --- /dev/null +++ b/examples/opf/finetune/FT2_operating_condition/config_HeteroHEAT_partial.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT2_operating_condition", + "_ft_description": "Operating-condition fine-tuning: fine-tune on a new distribution of load/generation profiles for pglib_opf_case14_ieee with a limited labeled budget (5000 samples).", + "_ft_case_name": "pglib_opf_case14_ieee", + "_ft_num_groups": "2", + "_ft_max_samples": 5000, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT2_operating_condition_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_full.json b/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_full.json new file mode 100644 index 000000000..d72c79a55 --- /dev/null +++ b/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_full.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT2_operating_condition", + "_ft_description": "Operating-condition fine-tuning: fine-tune on a new distribution of load/generation profiles for pglib_opf_case14_ieee with a limited labeled budget (5000 samples).", + "_ft_case_name": "pglib_opf_case14_ieee", + "_ft_num_groups": "2", + "_ft_max_samples": 5000, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT2_operating_condition_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_head_only.json b/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_head_only.json new file mode 100644 index 000000000..41c768db5 --- /dev/null +++ b/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_head_only.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT2_operating_condition", + "_ft_description": "Operating-condition fine-tuning: fine-tune on a new distribution of load/generation profiles for pglib_opf_case14_ieee with a limited labeled budget (5000 samples).", + "_ft_case_name": "pglib_opf_case14_ieee", + "_ft_num_groups": "2", + "_ft_max_samples": 5000, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT2_operating_condition_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_partial.json b/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_partial.json new file mode 100644 index 000000000..19be9a5b9 --- /dev/null +++ b/examples/opf/finetune/FT2_operating_condition/config_HeteroSAGE_partial.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT2_operating_condition", + "_ft_description": "Operating-condition fine-tuning: fine-tune on a new distribution of load/generation profiles for pglib_opf_case14_ieee with a limited labeled budget (5000 samples).", + "_ft_case_name": "pglib_opf_case14_ieee", + "_ft_num_groups": "2", + "_ft_max_samples": 5000, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT2_operating_condition_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT2_operating_condition/job-frontier-HeteroHEAT.sh b/examples/opf/finetune/FT2_operating_condition/job-frontier-HeteroHEAT.sh new file mode 100755 index 000000000..977b1d5f8 --- /dev/null +++ b/examples/opf/finetune/FT2_operating_condition/job-frontier-HeteroHEAT.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Convenience launcher for FT2_operating_condition with HeteroHEAT. +# Usage: sbatch FT2_operating_condition/job-frontier-HeteroHEAT.sh +# Override regime via SBATCH --export, e.g.: +# sbatch --export=ALL,FT_REGIME=partial FT2_operating_condition/job-frontier-HeteroHEAT.sh +#SBATCH -A LRN078 +#SBATCH -J OPF-FT2-HeteroHEAT +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT2_operating_condition-HeteroHEAT-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT2_operating_condition-HeteroHEAT-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +export FT_STRATEGY=FT2_operating_condition +export FT_ARCH=HeteroHEAT +export FT_REGIME=${FT_REGIME:-full} +export PRETRAINED_MODEL=HeteroHEAT_best + +bash $(dirname $0)/../job-frontier-finetune.sh diff --git a/examples/opf/finetune/FT2_operating_condition/job-frontier-HeteroSAGE.sh b/examples/opf/finetune/FT2_operating_condition/job-frontier-HeteroSAGE.sh new file mode 100755 index 000000000..fcfc8ba0b --- /dev/null +++ b/examples/opf/finetune/FT2_operating_condition/job-frontier-HeteroSAGE.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Convenience launcher for FT2_operating_condition with HeteroSAGE. +# Usage: sbatch FT2_operating_condition/job-frontier-HeteroSAGE.sh +# Override regime via SBATCH --export, e.g.: +# sbatch --export=ALL,FT_REGIME=partial FT2_operating_condition/job-frontier-HeteroSAGE.sh +#SBATCH -A LRN078 +#SBATCH -J OPF-FT2-HeteroSAGE +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT2_operating_condition-HeteroSAGE-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT2_operating_condition-HeteroSAGE-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +export FT_STRATEGY=FT2_operating_condition +export FT_ARCH=HeteroSAGE +export FT_REGIME=${FT_REGIME:-full} +export PRETRAINED_MODEL=HeteroSAGE_best + +bash $(dirname $0)/../job-frontier-finetune.sh diff --git a/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_full.json b/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_full.json new file mode 100644 index 000000000..5e359699a --- /dev/null +++ b/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_full.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT3_contingency", + "_ft_description": "Contingency fine-tuning: adapt to N-1 topological perturbations (topological_perturbations=True) for pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": true, + "ft_data_modelname": "FT3_contingency_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_head_only.json b/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_head_only.json new file mode 100644 index 000000000..69a141272 --- /dev/null +++ b/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_head_only.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT3_contingency", + "_ft_description": "Contingency fine-tuning: adapt to N-1 topological perturbations (topological_perturbations=True) for pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": true, + "ft_data_modelname": "FT3_contingency_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_partial.json b/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_partial.json new file mode 100644 index 000000000..e0210f576 --- /dev/null +++ b/examples/opf/finetune/FT3_contingency/config_HeteroHEAT_partial.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT3_contingency", + "_ft_description": "Contingency fine-tuning: adapt to N-1 topological perturbations (topological_perturbations=True) for pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": true, + "ft_data_modelname": "FT3_contingency_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_full.json b/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_full.json new file mode 100644 index 000000000..0cb28a810 --- /dev/null +++ b/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_full.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT3_contingency", + "_ft_description": "Contingency fine-tuning: adapt to N-1 topological perturbations (topological_perturbations=True) for pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": true, + "ft_data_modelname": "FT3_contingency_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_head_only.json b/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_head_only.json new file mode 100644 index 000000000..a18fbc38a --- /dev/null +++ b/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_head_only.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT3_contingency", + "_ft_description": "Contingency fine-tuning: adapt to N-1 topological perturbations (topological_perturbations=True) for pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": true, + "ft_data_modelname": "FT3_contingency_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_partial.json b/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_partial.json new file mode 100644 index 000000000..5c1db74c8 --- /dev/null +++ b/examples/opf/finetune/FT3_contingency/config_HeteroSAGE_partial.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT3_contingency", + "_ft_description": "Contingency fine-tuning: adapt to N-1 topological perturbations (topological_perturbations=True) for pglib_opf_case118_ieee.", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": true, + "ft_data_modelname": "FT3_contingency_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT3_contingency/job-frontier-HeteroHEAT.sh b/examples/opf/finetune/FT3_contingency/job-frontier-HeteroHEAT.sh new file mode 100755 index 000000000..242e295c5 --- /dev/null +++ b/examples/opf/finetune/FT3_contingency/job-frontier-HeteroHEAT.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Convenience launcher for FT3_contingency with HeteroHEAT. +# Usage: sbatch FT3_contingency/job-frontier-HeteroHEAT.sh +# Override regime via SBATCH --export, e.g.: +# sbatch --export=ALL,FT_REGIME=partial FT3_contingency/job-frontier-HeteroHEAT.sh +#SBATCH -A LRN078 +#SBATCH -J OPF-FT3-HeteroHEAT +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT3_contingency-HeteroHEAT-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT3_contingency-HeteroHEAT-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +export FT_STRATEGY=FT3_contingency +export FT_ARCH=HeteroHEAT +export FT_REGIME=${FT_REGIME:-full} +export PRETRAINED_MODEL=HeteroHEAT_best + +bash $(dirname $0)/../job-frontier-finetune.sh diff --git a/examples/opf/finetune/FT3_contingency/job-frontier-HeteroSAGE.sh b/examples/opf/finetune/FT3_contingency/job-frontier-HeteroSAGE.sh new file mode 100755 index 000000000..5851a5241 --- /dev/null +++ b/examples/opf/finetune/FT3_contingency/job-frontier-HeteroSAGE.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Convenience launcher for FT3_contingency with HeteroSAGE. +# Usage: sbatch FT3_contingency/job-frontier-HeteroSAGE.sh +# Override regime via SBATCH --export, e.g.: +# sbatch --export=ALL,FT_REGIME=partial FT3_contingency/job-frontier-HeteroSAGE.sh +#SBATCH -A LRN078 +#SBATCH -J OPF-FT3-HeteroSAGE +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT3_contingency-HeteroSAGE-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT3_contingency-HeteroSAGE-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +export FT_STRATEGY=FT3_contingency +export FT_ARCH=HeteroSAGE +export FT_REGIME=${FT_REGIME:-full} +export PRETRAINED_MODEL=HeteroSAGE_best + +bash $(dirname $0)/../job-frontier-finetune.sh diff --git a/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_full.json b/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_full.json new file mode 100644 index 000000000..fcb61ed15 --- /dev/null +++ b/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_full.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT4_task_specific", + "_ft_description": "Task-specific fine-tuning: adapt the pretrained bus-voltage encoder to generator dispatch prediction (node_target_type=generator).", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT4_task_specific_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "generator", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "generator_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_head_only.json b/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_head_only.json new file mode 100644 index 000000000..b6ade40cb --- /dev/null +++ b/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_head_only.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT4_task_specific", + "_ft_description": "Task-specific fine-tuning: adapt the pretrained bus-voltage encoder to generator dispatch prediction (node_target_type=generator).", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT4_task_specific_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "generator", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "generator_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_partial.json b/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_partial.json new file mode 100644 index 000000000..b3777139f --- /dev/null +++ b/examples/opf/finetune/FT4_task_specific/config_HeteroHEAT_partial.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT4_task_specific", + "_ft_description": "Task-specific fine-tuning: adapt the pretrained bus-voltage encoder to generator dispatch prediction (node_target_type=generator).", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT4_task_specific_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "generator", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "generator_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_full.json b/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_full.json new file mode 100644 index 000000000..15e58f483 --- /dev/null +++ b/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_full.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT4_task_specific", + "_ft_description": "Task-specific fine-tuning: adapt the pretrained bus-voltage encoder to generator dispatch prediction (node_target_type=generator).", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT4_task_specific_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "generator", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "generator_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "full" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_head_only.json b/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_head_only.json new file mode 100644 index 000000000..89be0d4eb --- /dev/null +++ b/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_head_only.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT4_task_specific", + "_ft_description": "Task-specific fine-tuning: adapt the pretrained bus-voltage encoder to generator dispatch prediction (node_target_type=generator).", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT4_task_specific_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "generator", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": true + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "generator_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "head_only" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_partial.json b/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_partial.json new file mode 100644 index 000000000..82b4a603c --- /dev/null +++ b/examples/opf/finetune/FT4_task_specific/config_HeteroSAGE_partial.json @@ -0,0 +1,137 @@ +{ + "_ft_strategy": "FT4_task_specific", + "_ft_description": "Task-specific fine-tuning: adapt the pretrained bus-voltage encoder to generator dispatch prediction (node_target_type=generator).", + "_ft_case_name": "pglib_opf_case118_ieee", + "_ft_num_groups": "20", + "_ft_max_samples": null, + "_ft_topological_perturbations": false, + "ft_data_modelname": "FT4_task_specific_data", + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "generator", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false, + "freeze_conv_layers": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "generator_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 15, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": false, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0005 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": "partial" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/finetune/FT4_task_specific/job-frontier-HeteroHEAT.sh b/examples/opf/finetune/FT4_task_specific/job-frontier-HeteroHEAT.sh new file mode 100755 index 000000000..6ab045ac3 --- /dev/null +++ b/examples/opf/finetune/FT4_task_specific/job-frontier-HeteroHEAT.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Convenience launcher for FT4_task_specific with HeteroHEAT. +# Usage: sbatch FT4_task_specific/job-frontier-HeteroHEAT.sh +# Override regime via SBATCH --export, e.g.: +# sbatch --export=ALL,FT_REGIME=partial FT4_task_specific/job-frontier-HeteroHEAT.sh +#SBATCH -A LRN078 +#SBATCH -J OPF-FT4-HeteroHEAT +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT4_task_specific-HeteroHEAT-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT4_task_specific-HeteroHEAT-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +export FT_STRATEGY=FT4_task_specific +export FT_ARCH=HeteroHEAT +export FT_REGIME=${FT_REGIME:-full} +export PRETRAINED_MODEL=HeteroHEAT_best + +bash $(dirname $0)/../job-frontier-finetune.sh diff --git a/examples/opf/finetune/FT4_task_specific/job-frontier-HeteroSAGE.sh b/examples/opf/finetune/FT4_task_specific/job-frontier-HeteroSAGE.sh new file mode 100755 index 000000000..b79087c17 --- /dev/null +++ b/examples/opf/finetune/FT4_task_specific/job-frontier-HeteroSAGE.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Convenience launcher for FT4_task_specific with HeteroSAGE. +# Usage: sbatch FT4_task_specific/job-frontier-HeteroSAGE.sh +# Override regime via SBATCH --export, e.g.: +# sbatch --export=ALL,FT_REGIME=partial FT4_task_specific/job-frontier-HeteroSAGE.sh +#SBATCH -A LRN078 +#SBATCH -J OPF-FT4-HeteroSAGE +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT4_task_specific-HeteroSAGE-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT4_task_specific-HeteroSAGE-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +export FT_STRATEGY=FT4_task_specific +export FT_ARCH=HeteroSAGE +export FT_REGIME=${FT_REGIME:-full} +export PRETRAINED_MODEL=HeteroSAGE_best + +bash $(dirname $0)/../job-frontier-finetune.sh diff --git a/examples/opf/finetune/collect_results.py b/examples/opf/finetune/collect_results.py new file mode 100644 index 000000000..15a35780f --- /dev/null +++ b/examples/opf/finetune/collect_results.py @@ -0,0 +1,282 @@ +"""Aggregate results.json files from all FT1 and FT3 experiment runs. + +Scans the logs/ directory for completed runs, reads their results.json and +training_curve.csv files, and writes two summary artefacts to results/: + results/ft1_ft3_summary.csv — one row per run, flat columns + results/ft1_ft3_summary.json — same data as structured JSON (for plots) + +Usage:: + + # From examples/opf/finetune/ + python collect_results.py + python collect_results.py --logs_root ../../logs --out_dir results + +The output files are intentionally human-readable and can be reloaded by +plot_ft_results.py to regenerate all figures without re-running experiments. +""" + +import argparse +import csv +import glob +import json +import os +import sys + + +# ───────────────────────────────────────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────────────────────────────────────── + +def _load_results_json(path: str) -> dict: + with open(path) as fh: + return json.load(fh) + + +def _load_training_curve(csv_path: str) -> list[dict]: + """Return list of {step, tag, value} rows.""" + if not os.path.isfile(csv_path): + return [] + rows = [] + with open(csv_path) as fh: + reader = csv.DictReader(fh) + for row in reader: + try: + rows.append({ + "step": int(row["step"]) if row["step"] else None, + "tag": row["tag"], + "value": float(row["value"]) if row["value"] else None, + }) + except (KeyError, ValueError): + continue + return rows + + +def _pivot_training_curve(rows: list[dict]) -> dict[str, list]: + """Convert flat rows to tag → list of (epoch, value) pairs.""" + from collections import defaultdict + result = defaultdict(list) + for r in rows: + if r["step"] is not None and r["value"] is not None: + result[r["tag"]].append((r["step"], r["value"])) + return dict(result) + + +# ───────────────────────────────────────────────────────────────────────────── +# Core collection logic +# ───────────────────────────────────────────────────────────────────────────── + +def collect(logs_root: str, out_dir: str) -> None: + """Scan logs_root for FT1/FT3 results and write summary files.""" + # Discover all results.json files for FT1 and FT3 runs. + # We match any log name that starts with FT1_ or FT3_. + patterns = [ + os.path.join(logs_root, "FT1_*", "results.json"), + os.path.join(logs_root, "FT3_*", "results.json"), + ] + found = [] + for pat in patterns: + found.extend(sorted(glob.glob(pat))) + + if not found: + print(f"[collect] No results.json files found under {logs_root}/FT1_* or FT3_*.") + print(" Have all experiments finished? Check that save_run_results() completed.") + sys.exit(0) + + print(f"[collect] Found {len(found)} results file(s).") + + records = [] # flat dicts for CSV + detailed = [] # full JSON payloads (with training curves) for JSON + + for results_path in found: + log_dir = os.path.dirname(results_path) + log_name = os.path.basename(log_dir) + + payload = _load_results_json(results_path) + meta = payload.get("meta", {}) + metrics = payload.get("test_metrics", {}) + + # Load training curve + curve_path = os.path.join(log_dir, "training_curve.csv") + curve_rows = _load_training_curve(curve_path) + curve = _pivot_training_curve(curve_rows) + + # ── Determine friendly run label ────────────────────────────────── + strategy = meta.get("ft_strategy", "unknown") + arch = meta.get("arch", "unknown") + regime = meta.get("regime", "full") + pretrained = meta.get("pretrained", True) + label = f"{arch}_{regime}" + ("" if pretrained else "_baseline") + + # ── Flat record (one row per run) ───────────────────────────────── + flat = { + "log_name": log_name, + "strategy": strategy, + "arch": arch, + "regime": regime, + "pretrained": pretrained, + "label": label, + # FT1 classification metrics + "bce": metrics.get("bce"), + "accuracy": metrics.get("accuracy"), + "precision": metrics.get("precision"), + "recall": metrics.get("recall"), + "f1": metrics.get("f1"), + "auc_roc": metrics.get("auc_roc"), + "n_samples": metrics.get("n_samples"), + # FT3 regression metrics + "overall_mse": metrics.get("overall_mse"), + "Va_mse": metrics.get("Va_mse"), + "Va_mae": metrics.get("Va_mae"), + "Va_r2": metrics.get("Va_r2"), + "Vm_mse": metrics.get("Vm_mse"), + "Vm_mae": metrics.get("Vm_mae"), + "Vm_r2": metrics.get("Vm_r2"), + "n_nodes": metrics.get("n_nodes"), + # Training hyperparameters + "num_epoch": meta.get("num_epoch"), + "learning_rate": meta.get("learning_rate"), + "config_file": meta.get("config_file"), + # Training curve extremes + "best_val_error": _best_val(curve), + "final_train_error": _final_train(curve), + "n_epochs_trained": _n_epochs(curve), + } + records.append(flat) + + # ── Detailed entry (full metrics + curve for plotting) ───────────── + detailed.append({ + "log_name": log_name, + "meta": meta, + "test_metrics": { + # Omit large arrays from the JSON summary; keep just scalars. + k: v for k, v in metrics.items() + if k not in ("probs", "labels", "preds_sample", "targets_sample") + }, + "probs": metrics.get("probs", []), + "labels": metrics.get("labels", []), + "preds_sample": metrics.get("preds_sample", []), + "targets_sample": metrics.get("targets_sample", []), + "training_curve": curve, + }) + + # ── Write outputs ────────────────────────────────────────────────────── + os.makedirs(out_dir, exist_ok=True) + + # CSV + csv_path = os.path.join(out_dir, "ft1_ft3_summary.csv") + if records: + fieldnames = list(records[0].keys()) + with open(csv_path, "w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=fieldnames) + writer.writeheader() + writer.writerows(records) + print(f"[collect] CSV → {csv_path} ({len(records)} rows)") + + # JSON + json_path = os.path.join(out_dir, "ft1_ft3_summary.json") + with open(json_path, "w") as fh: + json.dump(detailed, fh, indent=2, default=_json_default) + print(f"[collect] JSON → {json_path}") + + # ── Print quick summary table ────────────────────────────────────────── + _print_table(records) + + +# ───────────────────────────────────────────────────────────────────────────── +# Helpers +# ───────────────────────────────────────────────────────────────────────────── + +def _best_val(curve: dict) -> float | None: + vals = curve.get("validate error", []) + return min(v for _, v in vals) if vals else None + + +def _final_train(curve: dict) -> float | None: + vals = curve.get("train error", []) + return vals[-1][1] if vals else None + + +def _n_epochs(curve: dict) -> int | None: + vals = curve.get("train error", []) + return len(vals) if vals else None + + +def _json_default(obj): + import numpy as np + if isinstance(obj, (np.floating, np.integer)): + return obj.item() + if isinstance(obj, np.ndarray): + return obj.tolist() + return str(obj) + + +def _print_table(records: list[dict]) -> None: + """Print a compact ASCII summary table.""" + # Group by strategy + from itertools import groupby + for strategy, grp in groupby( + sorted(records, key=lambda r: (r["strategy"], r["arch"], r["regime"])), + key=lambda r: r["strategy"], + ): + rows = list(grp) + print(f"\n{'─'*72}") + print(f" {strategy}") + print(f"{'─'*72}") + is_ft1 = "FT1" in strategy + if is_ft1: + hdr = f" {'Label':<30} {'BCE':>7} {'Acc':>7} {'F1':>7} {'AUC':>7}" + else: + hdr = f" {'Label':<30} {'Va_MSE':>9} {'Vm_MSE':>9} {'Va_R2':>7} {'Vm_R2':>7}" + print(hdr) + for r in rows: + lbl = r["label"] + if is_ft1: + print( + f" {lbl:<30} " + f"{_fmt(r['bce']):>7} {_fmt(r['accuracy']):>7} " + f"{_fmt(r['f1']):>7} {_fmt(r['auc_roc']):>7}" + ) + else: + print( + f" {lbl:<30} " + f"{_fmt(r['Va_mse']):>9} {_fmt(r['Vm_mse']):>9} " + f"{_fmt(r['Va_r2']):>7} {_fmt(r['Vm_r2']):>7}" + ) + print(f"\n{'─'*72}\n") + + +def _fmt(v) -> str: + if v is None: + return " —" + try: + return f"{float(v):.4f}" + except (TypeError, ValueError): + return str(v) + + +# ───────────────────────────────────────────────────────────────────────────── +# CLI +# ───────────────────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser( + description="Aggregate FT1 and FT3 results from the logs/ directory." + ) + _ft_dir = os.path.dirname(os.path.abspath(__file__)) + _repo_root = os.path.normpath(os.path.join(_ft_dir, "..", "..", "..")) + parser.add_argument( + "--logs_root", + default=os.path.join(_repo_root, "logs"), + help="Root directory containing log subdirectories (default: /logs)", + ) + parser.add_argument( + "--out_dir", + default=os.path.join(_ft_dir, "results"), + help="Directory to write summary files (default: examples/opf/finetune/results/)", + ) + args = parser.parse_args() + collect(args.logs_root, args.out_dir) + + +if __name__ == "__main__": + main() diff --git a/examples/opf/finetune/ft_utils.py b/examples/opf/finetune/ft_utils.py new file mode 100644 index 000000000..934397667 --- /dev/null +++ b/examples/opf/finetune/ft_utils.py @@ -0,0 +1,373 @@ +"""Shared evaluation and statistics utilities for the FT1 and FT3 fine-tuning pipeline. + +Provides: + EpochCSVWriter — wraps a TensorBoard SummaryWriter, also writes scalars to CSV + evaluate_ft1 — classification metrics (accuracy, F1, AUC-ROC) gathered across ranks + evaluate_ft3 — regression metrics (per-dim MSE, MAE, R²) gathered across ranks + save_run_results — writes results.json to a run's log directory + load_best_or_last_checkpoint — loads the best available checkpoint for post-training eval +""" + +import csv +import glob +import json +import os + +import numpy as np +import torch +import torch.nn.functional as F + + +# ────────────────────────────────────────────────────────────────────────────── +# TensorBoard + CSV writer wrapper +# ────────────────────────────────────────────────────────────────────────────── + +class EpochCSVWriter: + """Wraps a TensorBoard SummaryWriter and *simultaneously* writes every + ``add_scalar`` call to a CSV file. + + The TB writer may be ``None`` (non-rank-0 processes) — CSV writing is + skipped in that case to keep files on rank 0 only. + + Usage:: + + tb = model_utils.get_summary_writer(log_name) # None on non-rank-0 + writer = EpochCSVWriter(tb, f"logs/{log_name}/training_curve.csv") + # Pass writer to train_validate_test as usual. + # At the end of the script call writer.close() — CSV is flushed then. + """ + + def __init__(self, tb_writer, csv_path: str): + self._tb = tb_writer + self._csv_path = csv_path + self._rows: list[dict] = [] # [{step, tag, value}] + + # ── intercept add_scalar ────────────────────────────────────────────── + + def add_scalar(self, tag, value, global_step=None, *args, **kwargs): + if self._tb is not None: + self._tb.add_scalar(tag, value, global_step, *args, **kwargs) + if self._tb is not None: # only rank-0 writes CSV + try: + v = float(value) + except (TypeError, ValueError): + v = None + self._rows.append({"step": global_step, "tag": tag, "value": v}) + + # ── forward everything else to the underlying writer ────────────────── + + def __getattr__(self, name): + # Called only for attributes not found on EpochCSVWriter itself + return getattr(self._tb, name) + + # ── close / flush ───────────────────────────────────────────────────── + + def close(self): + if self._tb is not None: + self._flush_csv() + self._tb.close() + + def _flush_csv(self): + if not self._rows: + return + os.makedirs(os.path.dirname(self._csv_path) or ".", exist_ok=True) + with open(self._csv_path, "w", newline="") as fh: + w = csv.DictWriter(fh, fieldnames=["step", "tag", "value"]) + w.writeheader() + w.writerows(self._rows) + + +# ────────────────────────────────────────────────────────────────────────────── +# Distributed evaluation helpers +# ────────────────────────────────────────────────────────────────────────────── + +def _gather_tensors(comm, local_tensor): + """MPI-gather a 2-D tensor from all ranks to rank 0. + + Returns the concatenated tensor on rank 0, None on all other ranks. + """ + all_tensors = comm.gather(local_tensor, root=0) + if comm.Get_rank() != 0: + return None + return torch.cat([t for t in all_tensors if t is not None and t.numel() > 0], dim=0) + + +def evaluate_ft1(model, test_loader, device, comm): + """Evaluate FT1 binary feasibility classification on the full test set. + + All ranks participate; metrics are computed and returned only on rank 0 + (other ranks receive ``None``). + + Parameters + ---------- + model : nn.Module (DDP-wrapped is fine) + test_loader : DataLoader yielding HeteroData with data.y ∈ {0.0, 1.0} + device : torch.device + comm : MPI communicator + + Returns (rank 0 only) + ------- + dict with keys: bce, accuracy, precision, recall, f1, auc_roc, + n_samples, n_feasible, n_infeasible, + probs (list), labels (list) + """ + model.eval() + local_logits: list[torch.Tensor] = [] + local_labels: list[torch.Tensor] = [] + + with torch.no_grad(): + for data in test_loader: + data = data.to(device) + pred = model(data) + # Handle (pred, pred_var) tuple from var_output models + if isinstance(pred, tuple): + pred = pred[0] + logits = pred[0] # [n_graphs, 1] + labels = data.y.float() + if labels.dim() == 1: + labels = labels.unsqueeze(1) + local_logits.append(logits.cpu().float()) + local_labels.append(labels.cpu().float()) + + if local_logits: + lc = torch.cat(local_logits, dim=0) # [local_N, 1] + ll = torch.cat(local_labels, dim=0) + else: + lc = torch.zeros(0, 1) + ll = torch.zeros(0, 1) + + logits = _gather_tensors(comm, lc) + labels = _gather_tensors(comm, ll) + + if comm.Get_rank() != 0: + return None + + from sklearn.metrics import ( + accuracy_score, f1_score, precision_score, + recall_score, roc_auc_score, + ) + + bce = F.binary_cross_entropy_with_logits(logits, labels).item() + probs_np = torch.sigmoid(logits).numpy().flatten() + labels_np = labels.numpy().flatten() + preds_bin = (probs_np >= 0.5).astype(int) + labels_int = labels_np.astype(int) + + has_both_classes = len(set(labels_int)) > 1 + + return { + "bce": bce, + "accuracy": float(accuracy_score(labels_int, preds_bin)), + "precision": float(precision_score(labels_int, preds_bin, zero_division=0)), + "recall": float(recall_score(labels_int, preds_bin, zero_division=0)), + "f1": float(f1_score(labels_int, preds_bin, zero_division=0)), + "auc_roc": float(roc_auc_score(labels_np, probs_np)) if has_both_classes else None, + "n_samples": int(len(labels_int)), + "n_feasible": int(labels_int.sum()), + "n_infeasible": int((1 - labels_int).sum()), + # Store raw arrays for ROC / confusion matrix plots + "probs": probs_np.tolist(), + "labels": labels_int.tolist(), + } + + +def evaluate_ft3(model, test_loader, device, comm, output_names=None): + """Evaluate FT3 node-level regression on the full test set. + + All ranks participate; metrics returned on rank 0 only (others get None). + + Parameters + ---------- + model : nn.Module + test_loader : NodeBatchAdapter-wrapped DataLoader + (each batch has data.y = [total_bus_nodes, out_dim]) + device : torch.device + comm : MPI communicator + output_names : list[str], default ["Va", "Vm"] + + Returns (rank 0 only) + ------- + dict with per-dimension keys: _mse, _mae, _r2, + plus overall_mse, n_nodes, preds (list), targets (list). + """ + if output_names is None: + output_names = ["Va", "Vm"] + + model.eval() + local_preds: list[torch.Tensor] = [] + local_targets: list[torch.Tensor] = [] + + with torch.no_grad(): + for data in test_loader: + data = data.to(device) + pred = model(data) + if isinstance(pred, tuple): + pred = pred[0] + node_pred = pred[0].cpu().float() # [total_nodes, out_dim] + node_true = data.y.cpu().float() + if node_true.dim() == 1: + node_true = node_true.unsqueeze(1) + # Guard against shape mismatches from padding + n = min(node_pred.shape[0], node_true.shape[0]) + local_preds.append(node_pred[:n]) + local_targets.append(node_true[:n]) + + if local_preds: + lp = torch.cat(local_preds, dim=0) + lt = torch.cat(local_targets, dim=0) + else: + out_dim = len(output_names) + lp = torch.zeros(0, out_dim) + lt = torch.zeros(0, out_dim) + + preds = _gather_tensors(comm, lp) + targets = _gather_tensors(comm, lt) + + if comm.Get_rank() != 0: + return None + + out_dim = min(preds.shape[1], len(output_names)) + metrics = { + "overall_mse": float(F.mse_loss(preds, targets).item()), + "n_nodes": int(preds.shape[0]), + } + for i, name in enumerate(output_names[:out_dim]): + p = preds[:, i] + t = targets[:, i] + mse = float(F.mse_loss(p, t).item()) + mae = float(F.l1_loss(p, t).item()) + ss_res = float(((t - p) ** 2).sum().item()) + ss_tot = float(((t - t.mean()) ** 2).sum().item()) + r2 = 1.0 - ss_res / (ss_tot + 1e-12) + metrics[f"{name}_mse"] = mse + metrics[f"{name}_mae"] = mae + metrics[f"{name}_r2"] = float(r2) + + # Store a subsample for scatter plots (max 2000 points per rank-0 data) + sample_idx = torch.randperm(preds.shape[0])[:2000] + metrics["preds_sample"] = preds[sample_idx].tolist() + metrics["targets_sample"] = targets[sample_idx].tolist() + + return metrics + + +# ────────────────────────────────────────────────────────────────────────────── +# Result persistence +# ────────────────────────────────────────────────────────────────────────────── + +def save_run_results(log_name: str, run_meta: dict, test_metrics: dict, + logs_root: str = "./logs"): + """Write ``results.json`` for a single run into its log directory. + + Parameters + ---------- + log_name : str — matches the training log name (``logs//``) + run_meta : dict — any descriptive fields: strategy, arch, regime, etc. + test_metrics : dict — output of evaluate_ft1 or evaluate_ft3 + logs_root : str — path to the logs directory (default ``"./logs"``) + """ + out_dir = os.path.join(logs_root, log_name) + os.makedirs(out_dir, exist_ok=True) + payload = {"log_name": log_name, "meta": run_meta, "test_metrics": test_metrics} + path = os.path.join(out_dir, "results.json") + with open(path, "w") as fh: + json.dump(payload, fh, indent=2, default=_json_default) + print(f"[results] Saved: {path}") + + +def _json_default(obj): + if isinstance(obj, (np.floating, np.integer)): + return obj.item() + if isinstance(obj, np.ndarray): + return obj.tolist() + raise TypeError(f"Object of type {type(obj)} is not JSON serializable") + + +# ────────────────────────────────────────────────────────────────────────────── +# Checkpoint loading +# ────────────────────────────────────────────────────────────────────────────── + +def load_best_or_last_checkpoint(log_name: str, model, logs_root: str = "./logs"): + """Load the best checkpoint from a run's log directory. + + HydraGNN's ``Checkpoint`` class saves checkpoints as ``.pk`` (a + symlink pointing to the epoch file with the best val_loss). The final + ``save_model`` call in the training script overwrites the symlink with the + last-epoch weights. + + This function preferentially loads the EPOCH checkpoint with the lowest + index under the log dir (the best checkpoint written by the Checkpoint + class before the final save overwrites the symlink). Falls back to + ``.pk`` if no epoch files exist. + + Returns True if a checkpoint was loaded, False otherwise. + """ + log_dir = os.path.join(logs_root, log_name) + if not os.path.isdir(log_dir): + return False + + # Prefer the best-epoch checkpoint (lowest val_loss, saved by Checkpoint) + # The symlink .pk always points to the LAST saved file. + # We can't reliably know the best epoch without re-reading the CSV, so we + # just load the symlink target (which IS the best if SaveCheckpoint fires + # and the final save_model hasn't been called yet — but after the script + # the symlink points to the last epoch). + # Best practical option: use the file with the SMALLEST val_loss based on + # training_curve.csv if available. + + csv_path = os.path.join(log_dir, "training_curve.csv") + best_epoch = _find_best_epoch_from_csv(csv_path) + + if best_epoch is not None: + cand = os.path.join(log_dir, f"{log_name}_epoch_{best_epoch}.pk") + if os.path.isfile(cand): + return _load_ckpt(cand, model) + + # Fallback: symlink / main checkpoint + main_ckpt = os.path.join(log_dir, f"{log_name}.pk") + if os.path.exists(main_ckpt): + return _load_ckpt(main_ckpt, model) + + # Fallback: any .pk file + pk_files = sorted(glob.glob(os.path.join(log_dir, "*.pk"))) + if pk_files: + return _load_ckpt(pk_files[-1], model) + + return False + + +def _find_best_epoch_from_csv(csv_path: str): + """Return epoch with lowest 'validate error' from training_curve.csv.""" + if not os.path.isfile(csv_path): + return None + try: + best_val, best_epoch = float("inf"), None + with open(csv_path) as fh: + reader = csv.DictReader(fh) + for row in reader: + if row.get("tag") == "validate error": + v = float(row["value"]) if row["value"] else float("inf") + if v < best_val: + best_val = v + best_epoch = row["step"] + return best_epoch + except Exception: + return None + + +def _load_ckpt(path: str, model) -> bool: + from collections import OrderedDict + try: + ckpt = torch.load(path, map_location="cpu") + sd = ckpt.get("model_state_dict", ckpt) + target = model.module if hasattr(model, "module") else model + own_keys = set(target.state_dict().keys()) + if own_keys and not next(iter(own_keys)).startswith("module"): + new_sd = OrderedDict((k.replace("module.", "", 1), v) for k, v in sd.items()) + sd = new_sd + target.load_state_dict(sd, strict=False) + print(f"[ft_utils] Loaded checkpoint: {path}") + return True + except Exception as exc: + print(f"[ft_utils] WARNING: could not load checkpoint {path}: {exc}") + return False diff --git a/examples/opf/finetune/generate_ft_configs.py b/examples/opf/finetune/generate_ft_configs.py new file mode 100644 index 000000000..68d576156 --- /dev/null +++ b/examples/opf/finetune/generate_ft_configs.py @@ -0,0 +1,326 @@ +"""Generate fine-tuning config JSON files for FT1–FT4 strategies. + +Run from the finetune/ directory: + python generate_ft_configs.py + +FT1 is a graph-level binary feasibility classification task: + - Feasible samples: taken from an existing OPF dataset (label = 1.0) + - Infeasible samples: synthesised by overloading load features (label = 0.0) + - Preprocessing: run generate_infeasible_samples.py before training FT1 + - Loss: binary_cross_entropy_with_logits (BCE) +""" +import json +import os + +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) + + +def _base_training(lr, num_epoch, regime): + return { + "num_epoch": num_epoch, + "batch_size": 32, + "patience": 15, + "early_stopping": True, + "Checkpoint": True, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": False, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [2, 3], + "voltage_output_index": -1, + }, + "Optimizer": {"type": "AdamW", "learning_rate": lr}, + "conv_checkpointing": False, + "loss_function_type": "mse", + "precision": "fp32", + "_ft_regime": regime, + } + + +def _base_arch(mpnn_type, hd, nl, freeze_conv, node_target_type, out_dim): + return { + "mpnn_type": mpnn_type, + "hidden_dim": hd, + "num_conv_layers": nl, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": {"ac_line": 9, "transformer": 11}, + "node_input_dims": {"bus": 4, "generator": 11, "load": 2, "shunt": 2}, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [32, 16], + "type": "mlp", + }, + } + ] + }, + "task_weights": [1.0], + "hetero_pooling_mode": "sum", + "node_target_type": node_target_type, + "global_attn_engine": None, + "global_attn_type": None, + "global_attn_heads": 0, + "output_dim": [out_dim], + "output_type": ["node"], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": None, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": False, + "freeze_conv_layers": freeze_conv, + } + + +def _base_arch_graph(mpnn_type, hd, nl, freeze_conv): + """Architecture config for graph-level binary classification (FT1).""" + return { + "mpnn_type": mpnn_type, + "hidden_dim": hd, + "num_conv_layers": nl, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": {"ac_line": 9, "transformer": 11}, + "node_input_dims": {"bus": 4, "generator": 11, "load": 2, "shunt": 2}, + "output_heads": { + "graph": [ + { + "type": "branch-0", + "architecture": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [32, 16], + }, + } + ] + }, + "task_weights": [1.0], + "hetero_pooling_mode": "sum", + "global_attn_engine": None, + "global_attn_type": None, + "global_attn_heads": 0, + "output_dim": [1], + "output_type": ["graph"], + # num_nodes is not fixed for the classification task (variable topology) + "num_nodes": None, + "input_dim": 4, + "pna_deg": None, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": False, + "freeze_conv_layers": freeze_conv, + } + + +VOI_BUS = { + "input_node_features": [0, 1, 2, 3], + "graph_feature_names": ["context"], + "graph_feature_dims": [1], + "node_feature_names": ["node_features"], + "node_feature_dims": [4], + "output_names": ["bus_solution"], + "output_index": [0], + "output_dim": [2], + "type": ["node"], + "denormalize_output": False, +} + +VOI_GEN = dict(VOI_BUS) +VOI_GEN["output_names"] = ["generator_solution"] +VOI_GEN["output_dim"] = [2] # generator solution = [Pg, Qg] (active/reactive dispatch) + +VOI_FEASIBILITY = { + "input_node_features": [0, 1, 2, 3], + "output_names": ["feasibility"], + "output_index": [0], + "output_dim": [1], + "type": ["graph"], + "denormalize_output": False, +} + +# Best HPO hyperparameters from Table VII of the manuscript +ARCHS = { + "HeteroSAGE": {"hd": 141, "nl": 5}, + "HeteroHEAT": {"hd": 232, "nl": 6}, +} + +# Learning rates per regime: +# head_only -> higher LR (only head updates; encoder frozen) +# partial -> medium LR (last conv + head) +# full -> lower LR (preserve pretrained representations) +REGIME_LR = { + "head_only": 1.0e-3, + "partial": 5.0e-4, + "full": 1.0e-4, +} + +STRATEGIES = { + "FT1_feasibility_classification": { + "desc": ( + "Feasibility classification: binary graph-level prediction of " + "whether an OPF instance is feasible (1) or infeasible (0). " + "Infeasible samples are synthesised by scaling load features by " + "an overload factor so that total demand exceeds generation capacity. " + "Preprocessing: run generate_infeasible_samples.py to create the " + "mixed dataset before training." + ), + # overload_factor is metadata only; actual value set in generate script + "overload_factor": 6.0, + "epochs": 50, + "target": "graph", + }, + "FT2_operating_condition": { + "desc": ( + "Operating-condition fine-tuning: fine-tune on a new distribution " + "of load/generation profiles for pglib_opf_case14_ieee with a " + "limited labeled budget (5000 samples)." + ), + "case": "pglib_opf_case14_ieee", + "groups": "2", + "max_samples": 5000, + "topo_perturb": False, + "epochs": 50, + "target": "bus", + }, + "FT3_contingency": { + "desc": ( + "Contingency fine-tuning: adapt to N-1 topological perturbations " + "(topological_perturbations=True) for pglib_opf_case118_ieee." + ), + "case": "pglib_opf_case118_ieee", + "groups": "20", + "max_samples": None, + "topo_perturb": True, + "epochs": 50, + "target": "bus", + }, + "FT4_task_specific": { + "desc": ( + "Task-specific fine-tuning: adapt the pretrained bus-voltage encoder " + "to generator dispatch prediction (node_target_type=generator)." + ), + "case": "pglib_opf_case118_ieee", + "groups": "20", + "max_samples": None, + "topo_perturb": False, + "epochs": 50, + "target": "generator", + }, +} + + +def _base_training_classify(lr, num_epoch, regime): + """Training config for binary classification tasks (FT1).""" + return { + "num_epoch": num_epoch, + "batch_size": 32, + "patience": 15, + "early_stopping": True, + "Checkpoint": True, + "checkpoint_warmup": 2, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": {"type": "AdamW", "learning_rate": lr}, + "conv_checkpointing": False, + "loss_function_type": "binary_cross_entropy", + "precision": "fp32", + "_ft_regime": regime, + } + + +def generate_all(): + for ft_dir, fm in STRATEGIES.items(): + out_dir = os.path.join(SCRIPT_DIR, ft_dir) + os.makedirs(out_dir, exist_ok=True) + + tgt = fm["target"] + + for arch_name, ap in ARCHS.items(): + for regime, freeze_conv in [ + ("head_only", True), + ("partial", False), + ("full", False), + ]: + lr = REGIME_LR[regime] + + if tgt == "graph": + # FT1: graph-level binary classification + arch = _base_arch_graph( + arch_name, ap["hd"], ap["nl"], freeze_conv + ) + training = _base_training_classify(lr, fm["epochs"], regime) + voi = VOI_FEASIBILITY + # Shared dataset (arch-independent) + data_modelname = "FT1_feasibility_data" + cfg = { + "_ft_strategy": ft_dir, + "_ft_description": fm["desc"], + "_ft_overload_factor": fm["overload_factor"], + "ft_data_modelname": data_modelname, + "Verbosity": {"level": 2}, + "NeuralNetwork": { + "Architecture": arch, + "Variables_of_interest": voi, + "Training": training, + }, + "Visualization": { + "plot_init_solution": False, + "plot_hist_solution": False, + "create_plots": False, + }, + } + else: + # FT2 / FT3 / FT4: node-level regression + out_dim = 2 # bus [Va, Vm] or generator [Pg, Qg] + node_target_type = tgt # "bus" or "generator" + voi = VOI_GEN if tgt == "generator" else VOI_BUS + arch = _base_arch( + arch_name, ap["hd"], ap["nl"], freeze_conv, node_target_type, out_dim + ) + training = _base_training(lr, fm["epochs"], regime) + cfg = { + "_ft_strategy": ft_dir, + "_ft_description": fm["desc"], + "_ft_case_name": fm["case"], + "_ft_num_groups": fm["groups"], + "_ft_max_samples": fm["max_samples"], + "_ft_topological_perturbations": fm["topo_perturb"], + "ft_data_modelname": f"{ft_dir}_{arch_name}_data", + "Verbosity": {"level": 2}, + "NeuralNetwork": { + "Architecture": arch, + "Variables_of_interest": voi, + "Training": training, + }, + "Visualization": { + "plot_init_solution": False, + "plot_hist_solution": False, + "create_plots": False, + }, + } + + fname = os.path.join(out_dir, f"config_{arch_name}_{regime}.json") + with open(fname, "w") as f: + json.dump(cfg, f, indent=4) + print(f"Created {fname}") + + print("\nAll fine-tuning configs generated.") + + +if __name__ == "__main__": + generate_all() diff --git a/examples/opf/finetune/generate_infeasible_samples.py b/examples/opf/finetune/generate_infeasible_samples.py new file mode 100644 index 000000000..5cd90709b --- /dev/null +++ b/examples/opf/finetune/generate_infeasible_samples.py @@ -0,0 +1,243 @@ +"""Generate FT1 feasibility-classification dataset from existing OPF HDF5 data. + +Infeasible OPF instances are synthesised by scaling all load features (Pd, Qd) +by a large overload factor, so that total demand exceeds total generation +capacity — guaranteeing AC-OPF infeasibility without running a solver. + +The output is a balanced HDF5 dataset (50 % feasible, 50 % infeasible) with a +graph-level binary label stored as ``data.y = torch.tensor([1.0])`` (feasible) +or ``data.y = torch.tensor([0.0])`` (infeasible). + +Usage (single rank, no MPI required):: + + python generate_infeasible_samples.py \\ + --src_dir ../dataset/FT3_contingency_HeteroSAGE_data.h5 \\ + --out_dir ../dataset/FT1_feasibility_data.h5 \\ + --overload_factor 6.0 \\ + --max_samples 5000 + +The output directory can then be used directly by train_opf_ft1_classify.py +via the config key ``"ft_data_modelname": "FT1_feasibility_data"``. +""" + +import os +import sys +import copy +import random +import argparse + +import torch +import torch.distributed as dist +from mpi4py import MPI + +# Make examples/opf importable (for opf_solution_utils if ever needed) +_OPF_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..") +sys.path.insert(0, _OPF_DIR) + +from hydragnn.utils.datasets.hdf5dataset import HDF5Dataset, HDF5Writer + + +# --------------------------------------------------------------------------- +# Sample manipulation helpers +# --------------------------------------------------------------------------- + +def strip_node_targets(data): + """Remove node-level prediction targets and y_loc; keep all node features.""" + for node_type in list(data.node_types): + node_store = data[node_type] + if hasattr(node_store, "y"): + del node_store["y"] + if hasattr(data, "y_loc"): + del data["y_loc"] + # Wipe top-level y so we can set the graph-level label cleanly + if hasattr(data, "y"): + del data["y"] + return data + + +def make_infeasible(data, overload_factor: float): + """Return a deep copy of *data* with load features scaled by *overload_factor*. + + Only the ``load`` node-type features are scaled (indices 0 = Pd, 1 = Qd). + All other features (bus, generator, shunt, edge_attr) are unchanged. + """ + infeasible = copy.deepcopy(data) + if "load" in infeasible.node_types and infeasible["load"].x is not None: + infeasible["load"].x = infeasible["load"].x * overload_factor + return infeasible + + +def label_samples(samples, label: float): + """Assign graph-level feasibility label to a list of samples in-place.""" + y = torch.tensor([label], dtype=torch.float32) + for s in samples: + s.y = y.clone() + return samples + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description="Generate FT1 feasibility classification HDF5 dataset.", + ) + parser.add_argument( + "--src_dir", + type=str, + required=True, + help=( + "Source HDF5 directory containing feasible OPF samples " + "(e.g. ../dataset/FT3_contingency_HeteroSAGE_data.h5)." + ), + ) + parser.add_argument( + "--out_dir", + type=str, + required=True, + help="Output HDF5 directory for the FT1 mixed feasibility dataset.", + ) + parser.add_argument( + "--overload_factor", + type=float, + default=6.0, + help=( + "Factor by which load features (Pd, Qd) are multiplied to create " + "infeasible samples. A value >=5 reliably causes infeasibility for " + "typical pglib-opf test cases." + ), + ) + parser.add_argument( + "--max_samples", + type=int, + default=None, + help=( + "Maximum number of feasible samples to use (randomly subsampled). " + "The same number of infeasible samples is generated, giving a " + "perfectly balanced dataset of size 2 * max_samples." + ), + ) + parser.add_argument( + "--train_frac", + type=float, + default=0.70, + help="Fraction of the mixed dataset assigned to the train split.", + ) + parser.add_argument( + "--val_frac", + type=float, + default=0.15, + help="Fraction assigned to the val split (remainder goes to test).", + ) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + + # Only rank-0 does the work (single-process preprocessing) + if rank != 0: + comm.Barrier() + return + + random.seed(args.seed) + torch.manual_seed(args.seed) + + # ── Load all feasible samples from the source HDF5 ───────────────────── + if not os.path.isdir(args.src_dir): + raise FileNotFoundError( + f"Source HDF5 directory not found: '{args.src_dir}'. " + "Run the FT3 (or any other OPF) preprocessing step first to " + "generate a feasible-sample HDF5 dataset." + ) + + # ── Reservoir-sample up to max_samples to avoid loading full dataset ─── + reservoir = [] + n_seen = 0 + cap = args.max_samples # None means no cap + + for split in ("trainset", "valset", "testset"): + try: + ds = HDF5Dataset(args.src_dir, split) + for i in range(len(ds)): + sample = ds[i] + strip_node_targets(sample) + if cap is None or n_seen < cap: + reservoir.append(sample) + else: + j = random.randint(0, n_seen) + if j < cap: + reservoir[j] = sample + n_seen += 1 + if (n_seen % 5000) == 0: + print(f" Scanned {n_seen} samples, kept {len(reservoir)}...") + except Exception as exc: + print(f" Warning: could not load split '{split}' from {args.src_dir}: {exc}") + + all_feasible = reservoir + print(f"Loaded {len(all_feasible)} feasible samples from {args.src_dir} (scanned {n_seen} total)") + + if len(all_feasible) == 0: + raise RuntimeError("No feasible samples found in the source dataset.") + + n_base = len(all_feasible) + + # ── Build feasible samples (label = 1.0) ─────────────────────────────── + feasible = label_samples(all_feasible, label=1.0) + + # ── Build infeasible samples (label = 0.0) ───────────────────────────── + print( + f"Generating {n_base} infeasible samples " + f"(overload_factor={args.overload_factor})..." + ) + infeasible = [make_infeasible(s, args.overload_factor) for s in all_feasible] + infeasible = label_samples(infeasible, label=0.0) + + # ── Mix and shuffle ──────────────────────────────────────────────────── + mixed = feasible + infeasible + random.shuffle(mixed) + n_total = len(mixed) + print(f"Total mixed samples: {n_total} ({n_base} feasible + {n_base} infeasible)") + + # ── Split ────────────────────────────────────────────────────────────── + n_train = int(args.train_frac * n_total) + n_val = int(args.val_frac * n_total) + splits = { + "trainset": mixed[:n_train], + "valset": mixed[n_train : n_train + n_val], + "testset": mixed[n_train + n_val :], + } + for name, s in splits.items(): + label_counts = { + 0: sum(1 for x in s if x.y.item() < 0.5), + 1: sum(1 for x in s if x.y.item() >= 0.5), + } + print( + f" {name}: {len(s)} samples " + f"(feasible={label_counts[1]}, infeasible={label_counts[0]})" + ) + + # ── Write HDF5 ───────────────────────────────────────────────────────── + print(f"Writing dataset to {args.out_dir} ...") + # HydraGNN's iterate_tqdm calls dist.get_rank() unconditionally; initialise + # a single-rank CPU process group so it does not raise on a login node. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + world_size=1, + rank=0, + init_method="tcp://127.0.0.1:29500", + ) + writer = HDF5Writer(args.out_dir, comm=MPI.COMM_SELF) + for split_name, split_samples in splits.items(): + writer.add(split_name, split_samples) + writer.save() + print("Done.") + + comm.Barrier() + + +if __name__ == "__main__": + main() diff --git a/examples/opf/finetune/job-frontier-FT1-bundled.sh b/examples/opf/finetune/job-frontier-FT1-bundled.sh new file mode 100755 index 000000000..8e5e279b8 --- /dev/null +++ b/examples/opf/finetune/job-frontier-FT1-bundled.sh @@ -0,0 +1,140 @@ +#!/bin/bash +# ============================================================================= +# Frontier Slurm job — FT1 Feasibility Classification (bundled data-efficiency) +# +# Runs all sample sizes sequentially; within each sample size the 4 training +# methods (full / partial / head_only / scratch) run concurrently with srun. +# +# Required env vars (set via sbatch --export or environment): +# FT_ARCH HeteroSAGE | HeteroHEAT +# +# Optional env vars: +# N_PER_RUN nodes per method (default: 8) +# PRETRAINED_MODEL override pretrained model name +# SAMPLE_SIZES space-separated list (default: 100 500 1000 2500 5000 +# 10000 25000 50000 100000) +# +# Usage (from examples/opf/finetune/): +# sbatch --export=ALL,FT_ARCH=HeteroSAGE \ +# job-frontier-FT1-bundled.sh +# ============================================================================= +#SBATCH -A eng164 +#SBATCH -J FT1-bundled +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT1-bundled-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT1-bundled-%j.out +#SBATCH -t 12:00:00 +#SBATCH -p batch +#SBATCH -N 32 + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- +FT_ARCH=${FT_ARCH:-HeteroSAGE} +N_PER_RUN=${N_PER_RUN:-8} +PRETRAINED_MODEL=${PRETRAINED_MODEL:-${FT_ARCH}_best} +read -ra SAMPLE_SIZES <<< "${SAMPLE_SIZES:-100 500 1000 2500 5000 10000 25000 50000 100000}" + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:${PYTHONPATH:-} + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +# ----------------------------------------------------------------------------- +# Derived +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +DATA_ROOT=$FT_DIR/../dataset + +echo "============================================================" +echo " FT1 Feasibility Classification — bundled data-efficiency" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Methods : full / partial / head_only / scratch (concurrent)" +echo " Nodes per run : $N_PER_RUN (total: $((N_PER_RUN * 4)))" +echo " Sample sizes : ${SAMPLE_SIZES[*]}" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +# Helper: launch one srun in the background for a given regime and sample size +_launch() { + local REGIME="$1" + local MAX_TRAIN_SAMPLES="$2" + local SCRATCH_FLAG="${3:-}" + local _scratch_tag="" + [[ -n "$SCRATCH_FLAG" ]] && _scratch_tag="_scratch" + local LOG_NAME="FT1_feasibility_${FT_ARCH}_${REGIME}${_scratch_tag}_n${MAX_TRAIN_SAMPLES}" + local LOGFILE="$HYDRAGNN_ROOT/${LOG_NAME}-${SLURM_JOB_ID}.out" + + echo " Launching $LOG_NAME → $LOGFILE" + + srun --exact -N${N_PER_RUN} -n$((N_PER_RUN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_ft1_classify.py \ + --inputfile FT1_feasibility_classification/config_${FT_ARCH}_${REGIME}.json \ + --modelname "$LOG_NAME" \ + --data_root $DATA_ROOT \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $REGIME \ + $SCRATCH_FLAG \ + --max_train_samples $MAX_TRAIN_SAMPLES \ + > "$LOGFILE" 2>&1 & +} + +# ----------------------------------------------------------------------------- +# Main loop: iterate sample sizes sequentially, methods concurrently +# ----------------------------------------------------------------------------- +ROUND=0 +for N in "${SAMPLE_SIZES[@]}"; do + ROUND=$((ROUND + 1)) + TOTAL=$((N * 2)) # balanced: N feasible + N infeasible + echo "" + echo "── Round $ROUND / ${#SAMPLE_SIZES[@]}: N=$N (total train samples=$TOTAL) ──" + + _launch full "$TOTAL" + _launch partial "$TOTAL" + _launch head_only "$TOTAL" + _launch full "$TOTAL" "--no_pretrained" # scratch baseline + + echo " Waiting for round $ROUND to finish..." + wait + echo " Round $ROUND done." +done + +echo "" +echo "All rounds finished. Job ID: $SLURM_JOB_ID" diff --git a/examples/opf/finetune/job-frontier-FT1-single-method.sh b/examples/opf/finetune/job-frontier-FT1-single-method.sh new file mode 100755 index 000000000..98e0e899f --- /dev/null +++ b/examples/opf/finetune/job-frontier-FT1-single-method.sh @@ -0,0 +1,152 @@ +#!/bin/bash +# ============================================================================= +# Frontier Slurm job — FT1 Feasibility Classification (single method) +# +# Runs ONE training method for all sample sizes sequentially. +# Uses only N_PER_RUN nodes (default: 8) — 4x smaller than the bundled job. +# +# Required env vars: +# FT_ARCH HeteroSAGE | HeteroHEAT +# FT_METHOD full | partial | head_only | scratch +# +# Optional env vars: +# N_PER_RUN nodes (default: 8) +# PRETRAINED_MODEL +# SAMPLE_SIZES space-separated list +# (default: 100 500 1000 2500 5000 10000 25000 50000 100000) +# +# Usage: +# sbatch --export=ALL,FT_ARCH=HeteroSAGE,FT_METHOD=full \ +# job-frontier-FT1-single-method.sh +# ============================================================================= +#SBATCH -A eng164 +#SBATCH -J FT1-single +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT1-single-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT1-single-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 1 + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- +FT_ARCH=${FT_ARCH:-HeteroSAGE} +FT_METHOD=${FT_METHOD:-full} +N_PER_RUN=${N_PER_RUN:-1} +PRETRAINED_MODEL=${PRETRAINED_MODEL:-${FT_ARCH}_best} +read -ra SAMPLE_SIZES <<< "${SAMPLE_SIZES:-100 500 1000 2500 5000 10000 25000 50000 100000}" + +# scratch baseline uses full regime + --no_pretrained flag +SCRATCH_FLAG="" +REGIME="$FT_METHOD" +if [[ "$FT_METHOD" == "scratch" ]]; then + REGIME="full" + SCRATCH_FLAG="--no_pretrained" +fi + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:${PYTHONPATH:-} + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +export OMP_NUM_THREADS=7 +export HYDRAGNN_AGGR_BACKEND=mpi +export HYDRAGNN_VALTEST=1 +export MPICH_ENV_DISPLAY=0 +export MPICH_VERSION_DISPLAY=0 +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH=/tmp +export PYTHONNOUSERSITE=1 + +# Single-node NCCL: AWS OFI / libfabric cxi plugin cannot bind without a remote +# peer. Disable network plugin so NCCL uses intra-node xGMI/SHM only. +if [[ "${N_PER_RUN:-1}" -le 1 ]]; then + unset NCCL_NET_PLUGIN NCCL_NET NCCL_NET_GDR_LEVEL NCCL_CROSS_NIC NCCL_SOCKET_IFNAME + export NCCL_P2P_LEVEL=SYS +fi + +# ----------------------------------------------------------------------------- +# Derived +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +DATA_ROOT=$FT_DIR/../dataset +_scratch_tag="" +[[ -n "$SCRATCH_FLAG" ]] && _scratch_tag="_scratch" + +echo "============================================================" +echo " FT1 Feasibility Classification — single method" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Method : $FT_METHOD (regime: $REGIME)" +echo " Nodes : $N_PER_RUN" +echo " Sample sizes : ${SAMPLE_SIZES[*]}" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +ROUND=0 +for N in "${SAMPLE_SIZES[@]}"; do + ROUND=$((ROUND + 1)) + TOTAL=$((N * 2)) + # NOTE: LOG_NAME is keyed on the *requested* training-sample count N + # (not TOTAL=N*2 which also includes the held-out class sub-sample), + # so each sweep point gets its own dir and never overwrites another. + LOG_NAME="FT1_feasibility_${FT_ARCH}_${REGIME}${_scratch_tag}_n${N}" + LOGFILE="$HYDRAGNN_ROOT/${LOG_NAME}-${SLURM_JOB_ID}.out" + + echo "" + echo "── Round $ROUND / ${#SAMPLE_SIZES[@]}: N=$N (total=$TOTAL) → $LOGFILE" + + srun -N${N_PER_RUN} -n$((N_PER_RUN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_ft1_classify.py \ + --inputfile FT1_feasibility_classification/config_${FT_ARCH}_${REGIME}.json \ + --modelname "$LOG_NAME" \ + --data_root $DATA_ROOT \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $REGIME \ + $SCRATCH_FLAG \ + --max_train_samples $TOTAL \ + > "$LOGFILE" 2>&1 + + echo " Round $ROUND done." +done + +echo "" +echo "All rounds finished. Job ID: $SLURM_JOB_ID" diff --git a/examples/opf/finetune/job-frontier-FT3-bundled.sh b/examples/opf/finetune/job-frontier-FT3-bundled.sh new file mode 100755 index 000000000..abd2c3ad5 --- /dev/null +++ b/examples/opf/finetune/job-frontier-FT3-bundled.sh @@ -0,0 +1,142 @@ +#!/bin/bash +# ============================================================================= +# Frontier Slurm job — FT3 N-1 Contingency OPF Regression (bundled data-efficiency) +# +# Runs all sample sizes sequentially; within each sample size the 4 training +# methods (full / partial / head_only / scratch) run concurrently with srun. +# +# Required env vars (set via sbatch --export or environment): +# FT_ARCH HeteroSAGE | HeteroHEAT +# +# Optional env vars: +# N_PER_RUN nodes per method (default: 8) +# PRETRAINED_MODEL override pretrained model name +# SAMPLE_SIZES space-separated list (default: 100 500 1000 2500 5000 +# 10000 25000 50000 100000 270000) +# +# Usage (from examples/opf/finetune/): +# sbatch --export=ALL,FT_ARCH=HeteroSAGE \ +# job-frontier-FT3-bundled.sh +# ============================================================================= +#SBATCH -A eng164 +#SBATCH -J FT3-bundled +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT3-bundled-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT3-bundled-%j.out +#SBATCH -t 12:00:00 +#SBATCH -p batch +#SBATCH -N 32 + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- +FT_ARCH=${FT_ARCH:-HeteroSAGE} +N_PER_RUN=${N_PER_RUN:-8} +PRETRAINED_MODEL=${PRETRAINED_MODEL:-${FT_ARCH}_best} +FT_STRATEGY=FT3_contingency +read -ra SAMPLE_SIZES <<< "${SAMPLE_SIZES:-100 500 1000 2500 5000 10000 25000 50000 100000 270000}" + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:${PYTHONPATH:-} + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +# ----------------------------------------------------------------------------- +# Derived +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +DATA_ROOT=$FT_DIR/../dataset + +echo "============================================================" +echo " FT3 N-1 Contingency OPF — bundled data-efficiency" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Methods : full / partial / head_only / scratch (concurrent)" +echo " Nodes per run : $N_PER_RUN (total: $((N_PER_RUN * 4)))" +echo " Sample sizes : ${SAMPLE_SIZES[*]}" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +# Helper: launch one srun in the background for a given regime and sample size +_launch() { + local REGIME="$1" + local MAX_TRAIN_SAMPLES="$2" + local SCRATCH_FLAG="${3:-}" + local _scratch_tag="" + [[ -n "$SCRATCH_FLAG" ]] && _scratch_tag="_scratch" + local LOG_NAME="finetune_${FT_STRATEGY}_${FT_ARCH}_${REGIME}${_scratch_tag}_n${MAX_TRAIN_SAMPLES}" + local LOGFILE="$HYDRAGNN_ROOT/${LOG_NAME}-${SLURM_JOB_ID}.out" + + echo " Launching $LOG_NAME → $LOGFILE" + + srun --exact -N${N_PER_RUN} -n$((N_PER_RUN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_finetune.py \ + --inputfile ${FT_STRATEGY}/config_${FT_ARCH}_${REGIME}.json \ + --hdf5 \ + --modelname "$LOG_NAME" \ + --resume_if_exists \ + --data_root $DATA_ROOT \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $REGIME \ + $SCRATCH_FLAG \ + --max_train_samples $MAX_TRAIN_SAMPLES \ + > "$LOGFILE" 2>&1 & +} + +# ----------------------------------------------------------------------------- +# Main loop: iterate sample sizes sequentially, methods concurrently +# ----------------------------------------------------------------------------- +ROUND=0 +for N in "${SAMPLE_SIZES[@]}"; do + ROUND=$((ROUND + 1)) + echo "" + echo "── Round $ROUND / ${#SAMPLE_SIZES[@]}: N=$N training samples ──" + + _launch full "$N" + _launch partial "$N" + _launch head_only "$N" + _launch full "$N" "--no_pretrained" # scratch baseline + + echo " Waiting for round $ROUND to finish..." + wait + echo " Round $ROUND done." +done + +echo "" +echo "All rounds finished. Job ID: $SLURM_JOB_ID" diff --git a/examples/opf/finetune/job-frontier-FT3-multi.sh b/examples/opf/finetune/job-frontier-FT3-multi.sh new file mode 100644 index 000000000..1a2adf37a --- /dev/null +++ b/examples/opf/finetune/job-frontier-FT3-multi.sh @@ -0,0 +1,135 @@ +#!/bin/bash +# ============================================================================= +# Frontier Slurm job — FT3 N-1 Contingency OPF Regression (multi-method) +# +# Runs all 4 training methods in parallel within a single allocation: +# - full fine-tuning (all layers) +# - partial fine-tuning (last conv + head) +# - head-only fine-tuning (linear probe) +# - from-scratch baseline (random init) +# +# Each method gets N_PER_RUN nodes; total allocation = 4 × N_PER_RUN. +# +# Required env vars (set via sbatch --export or environment): +# FT_ARCH HeteroSAGE | HeteroHEAT +# +# Optional env vars: +# MAX_TRAIN_SAMPLES limit training samples (default: use all) +# N_PER_RUN nodes per method (default: 8) +# PRETRAINED_MODEL override pretrained model name +# +# Usage (from examples/opf/finetune/): +# sbatch --export=ALL,FT_ARCH=HeteroSAGE,MAX_TRAIN_SAMPLES=5000 \ +# job-frontier-FT3-multi.sh +# ============================================================================= +#SBATCH -A LRN070 +#SBATCH -J OPF-FT3-MULTI +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT3-multi-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT3-multi-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -N 32 + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- +FT_ARCH=${FT_ARCH:-HeteroSAGE} +MAX_TRAIN_SAMPLES=${MAX_TRAIN_SAMPLES:-} +N_PER_RUN=${N_PER_RUN:-8} +PRETRAINED_MODEL=${PRETRAINED_MODEL:-${FT_ARCH}_best} +FT_STRATEGY=FT3_contingency + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:${PYTHONPATH:-} + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +# ----------------------------------------------------------------------------- +# Derived +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +DATA_ROOT=$FT_DIR/../dataset +_n_tag=${MAX_TRAIN_SAMPLES:+_n${MAX_TRAIN_SAMPLES}} + +echo "============================================================" +echo " FT3 N-1 Contingency OPF — multi-method job" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Methods : full / partial / head_only / scratch (parallel)" +echo " Nodes per run : $N_PER_RUN (total: $((N_PER_RUN * 4)))" +echo " Max train samples: ${MAX_TRAIN_SAMPLES:-all}" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +# Helper: launch one srun in the background, redirect output to its own log +_launch() { + local REGIME="$1" + local SCRATCH_FLAG="${2:-}" + local _scratch_tag="" + [[ -n "$SCRATCH_FLAG" ]] && _scratch_tag="_scratch" + local LOG_NAME="finetune_${FT_STRATEGY}_${FT_ARCH}_${REGIME}${_scratch_tag}${_n_tag}" + local LOGFILE="$HYDRAGNN_ROOT/${LOG_NAME}-${SLURM_JOB_ID}.out" + + echo " Launching $LOG_NAME → $LOGFILE" + + srun --exact -N${N_PER_RUN} -n$((N_PER_RUN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_finetune.py \ + --inputfile ${FT_STRATEGY}/config_${FT_ARCH}_${REGIME}.json \ + --hdf5 \ + --modelname "$LOG_NAME" \ + --data_root $DATA_ROOT \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $REGIME \ + $SCRATCH_FLAG \ + ${MAX_TRAIN_SAMPLES:+--max_train_samples $MAX_TRAIN_SAMPLES} \ + > "$LOGFILE" 2>&1 & +} + +# Launch all 4 methods in parallel +_launch full +_launch partial +_launch head_only +_launch full "--no_pretrained" # scratch baseline + +echo "" +echo "All 4 methods launched. Waiting for completion..." +wait +echo "" +echo "All methods finished. Job ID: $SLURM_JOB_ID" diff --git a/examples/opf/finetune/job-frontier-FT3-single-method.sh b/examples/opf/finetune/job-frontier-FT3-single-method.sh new file mode 100755 index 000000000..42565740d --- /dev/null +++ b/examples/opf/finetune/job-frontier-FT3-single-method.sh @@ -0,0 +1,159 @@ +#!/bin/bash +set -euo pipefail +# ============================================================================= +# Frontier Slurm job — FT3 N-1 Contingency OPF Regression (single method) +# +# Runs ONE training method for all sample sizes sequentially. +# Uses only N_PER_RUN nodes (default: 8) — 4x smaller than the bundled job. +# +# Required env vars: +# FT_ARCH HeteroSAGE | HeteroHEAT +# FT_METHOD full | partial | head_only | scratch +# +# Optional env vars: +# N_PER_RUN nodes (default: 8) +# PRETRAINED_MODEL +# SAMPLE_SIZES space-separated list +# (default: 100 500 1000 2500 5000 10000 25000 50000 +# 100000 270000) +# +# Usage: +# sbatch --export=ALL,FT_ARCH=HeteroSAGE,FT_METHOD=full \ +# job-frontier-FT3-single-method.sh +# ============================================================================= +#SBATCH -A eng164 +#SBATCH -J FT3-single +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/FT3-single-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/FT3-single-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 1 + +# ----------------------------------------------------------------------------- +# Configuration +# ----------------------------------------------------------------------------- +FT_ARCH=${FT_ARCH:-HeteroSAGE} +FT_METHOD=${FT_METHOD:-full} +N_PER_RUN=${N_PER_RUN:-1} +PRETRAINED_MODEL=${PRETRAINED_MODEL:-${FT_ARCH}_best} +FT_STRATEGY=FT3_contingency +read -ra SAMPLE_SIZES <<< "${SAMPLE_SIZES:-100 500 1000 2500 5000 10000 25000 50000 100000 270000}" + +SCRATCH_FLAG="" +REGIME="$FT_METHOD" +if [[ "$FT_METHOD" == "scratch" ]]; then + REGIME="full" + SCRATCH_FLAG="--no_pretrained" +fi + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:${PYTHONPATH:-} + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +export OMP_NUM_THREADS=7 +export HYDRAGNN_AGGR_BACKEND=mpi +export HYDRAGNN_VALTEST=1 +export MPICH_ENV_DISPLAY=0 +export MPICH_VERSION_DISPLAY=0 +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH=/tmp +export PYTHONNOUSERSITE=1 + +# Single-node NCCL: AWS OFI / libfabric cxi plugin cannot bind without a remote +# peer. Disable network plugin so NCCL uses intra-node xGMI/SHM only. +if [[ "${N_PER_RUN:-1}" -le 1 ]]; then + unset NCCL_NET_PLUGIN NCCL_NET NCCL_NET_GDR_LEVEL NCCL_CROSS_NIC NCCL_SOCKET_IFNAME + export NCCL_P2P_LEVEL=SYS +fi + +# ----------------------------------------------------------------------------- +# Derived +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +DATA_ROOT=$FT_DIR/../dataset +_scratch_tag="" +[[ -n "$SCRATCH_FLAG" ]] && _scratch_tag="_scratch" + +echo "============================================================" +echo " FT3 N-1 Contingency OPF — single method" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Method : $FT_METHOD (regime: $REGIME)" +echo " Nodes : $N_PER_RUN" +echo " Sample sizes : ${SAMPLE_SIZES[*]}" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +ROUND=0 +for N in "${SAMPLE_SIZES[@]}"; do + ROUND=$((ROUND + 1)) + LOG_NAME="finetune_${FT_STRATEGY}_${FT_ARCH}_${REGIME}${_scratch_tag}_n${N}" + LOGFILE="$HYDRAGNN_ROOT/${LOG_NAME}-${SLURM_JOB_ID}.out" + + echo "" + echo "── Round $ROUND / ${#SAMPLE_SIZES[@]}: N=$N → $LOGFILE" + + srun -N${N_PER_RUN} -n$((N_PER_RUN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_finetune.py \ + --inputfile ${FT_STRATEGY}/config_${FT_ARCH}_${REGIME}.json \ + --hdf5 \ + --modelname "$LOG_NAME" \ + --resume_if_exists \ + --data_root $DATA_ROOT \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $REGIME \ + $SCRATCH_FLAG \ + --max_train_samples $N \ + > "$LOGFILE" 2>&1 + + RESULT_JSON="$FT_DIR/logs/$LOG_NAME/results.json" + if [[ ! -f "$RESULT_JSON" ]]; then + echo "ERROR: Missing expected result artifact: $RESULT_JSON" + echo "See round log: $LOGFILE" + exit 1 + fi + + echo " Round $ROUND done." +done + +echo "" +echo "All rounds finished. Job ID: $SLURM_JOB_ID" diff --git a/examples/opf/finetune/job-frontier-finetune.sh b/examples/opf/finetune/job-frontier-finetune.sh new file mode 100644 index 000000000..76b9b44dd --- /dev/null +++ b/examples/opf/finetune/job-frontier-finetune.sh @@ -0,0 +1,171 @@ +#!/bin/bash +# ============================================================================= +# Frontier fine-tuning job template +# +# Usage: copy or symlink this file into any FT strategy directory, then edit +# the FT_STRATEGY and FT_ARCH_REGIME variables below, or override from CLI: +# +# sbatch --export=ALL,FT_STRATEGY=FT1_topology,FT_ARCH=HeteroSAGE,FT_REGIME=full \ +# job-frontier-finetune.sh +# +# Two-phase workflow: +# Phase 1 (--preonly) : preprocessing/serialisation with N_PREONLY nodes +# Phase 2 (training) : fine-tuning with N_TRAIN nodes +# +# To run both phases in one job, set PHASES=preonly,train (default). +# To run only one phase: PHASES=preonly or PHASES=train +# ============================================================================= +#SBATCH -A LRN078 +#SBATCH -J OPF-FT +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/finetune-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/finetune-%j.out +#SBATCH -t 04:00:00 +#SBATCH -p batch +#SBATCH -N 8 + +# ----------------------------------------------------------------------------- +# Fine-tuning configuration — edit these or override via --export +# ----------------------------------------------------------------------------- +FT_STRATEGY=${FT_STRATEGY:-FT2_operating_condition} # FT2_operating_condition | FT3_contingency | FT4_task_specific (use job-frontier-FT1-*.sh for FT1) +FT_ARCH=${FT_ARCH:-HeteroSAGE} # HeteroSAGE | HeteroHEAT +FT_REGIME=${FT_REGIME:-full} # full | partial | head_only +PRETRAINED_MODEL=${PRETRAINED_MODEL:-${FT_ARCH}_best} +PHASES=${PHASES:-preonly,train} # preonly,train | preonly | train +# Set NO_PRETRAINED=1 to skip loading pretrained weights (baseline comparison) +NO_PRETRAINED=${NO_PRETRAINED:-0} + +# Preprocessing node count (smaller than training — single node is enough for +# fine-tuning datasets which are much smaller than the 3M pretraining corpus). +N_PREONLY=1 +# Training node count +N_TRAIN=${SLURM_JOB_NUM_NODES} + +# Optional: limit labeled samples for the fine-tuning split +MAX_SAMPLES=${MAX_SAMPLES:-} # leave empty to use all samples + +# ----------------------------------------------------------------------------- +# Environment +# ----------------------------------------------------------------------------- +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +module unload darshan-runtime + +# ROCm / libfabric tuning (same as pretraining jobs) +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +# ----------------------------------------------------------------------------- +# Derived paths and names +# ----------------------------------------------------------------------------- +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune +INPUTFILE=$FT_DIR/${FT_STRATEGY}/config_${FT_ARCH}_${FT_REGIME}.json +_n_tag=${MAX_TRAIN_SAMPLES:+_n${MAX_TRAIN_SAMPLES}} +_scratch_tag=${NO_PRETRAINED:+_scratch} +LOG_NAME="finetune_${FT_STRATEGY}_${FT_ARCH}_${FT_REGIME}${_scratch_tag}${_n_tag}" +DATA_MODELNAME="${FT_STRATEGY}_data" + +# Read case/group metadata from the config (informational only) +CASE_NAME=$(python3 -c "import json; c=json.load(open('$INPUTFILE')); print(c.get('_ft_case_name',''))" 2>/dev/null) +NUM_GROUPS=$(python3 -c "import json; c=json.load(open('$INPUTFILE')); print(c.get('_ft_num_groups','1'))" 2>/dev/null) +TOPO_PERTURB=$(python3 -c "import json; c=json.load(open('$INPUTFILE')); print('--topological_perturbations' if c.get('_ft_topological_perturbations') else '')" 2>/dev/null) + +echo "============================================================" +echo " OPF Fine-tuning on Frontier" +echo " Strategy : $FT_STRATEGY" +echo " Arch : $FT_ARCH (pretrained: $PRETRAINED_MODEL)" +echo " Regime : $FT_REGIME" +echo " Phases : $PHASES" +echo " Case : $CASE_NAME (groups: $NUM_GROUPS)" +echo " Log name : $LOG_NAME" +echo " Job ID : $SLURM_JOB_ID" +echo "============================================================" + +cd $FT_DIR + +# Build optional --max_samples flag +MAX_SAMPLES_FLAG="" +if [ -n "$MAX_SAMPLES" ]; then + MAX_SAMPLES_FLAG="--max_samples $MAX_SAMPLES" +fi + +# Also read max_samples from config if not overridden from environment +if [ -z "$MAX_SAMPLES_FLAG" ]; then + CFG_MAX=$(python3 -c "import json; c=json.load(open('$INPUTFILE')); v=c.get('_ft_max_samples'); print(v if v is not None else '')" 2>/dev/null) + if [ -n "$CFG_MAX" ]; then + MAX_SAMPLES_FLAG="--max_samples $CFG_MAX" + fi +fi + +# ============================================================================= +# Phase 1: Preprocessing / serialisation +# ============================================================================= +if [[ "$PHASES" == *"preonly"* ]]; then + echo "" + echo "--- Phase 1: Preprocessing (--preonly) ---" + srun -N$N_PREONLY -n$((N_PREONLY * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u $HYDRAGNN_ROOT/examples/opf/train_opf_solution_heterogeneous.py \ + --inputfile $INPUTFILE \ + --hdf5 \ + --preonly \ + --case_name $CASE_NAME \ + --num_groups $NUM_GROUPS \ + --modelname $DATA_MODELNAME \ + --data_root $FT_DIR/../dataset \ + $TOPO_PERTURB \ + $MAX_SAMPLES_FLAG + echo "--- Phase 1 complete ---" +fi + +# ============================================================================= +# Phase 2: Fine-tuning +# ============================================================================= +if [[ "$PHASES" == *"train"* ]]; then + echo "" + echo "--- Phase 2: Fine-tuning ($FT_REGIME) with $FT_ARCH ---" + srun -N$N_TRAIN -n$((N_TRAIN * 8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_finetune.py \ + --inputfile ${FT_STRATEGY}/config_${FT_ARCH}_${FT_REGIME}.json \ + --hdf5 \ + --modelname $LOG_NAME \ + --data_root $FT_DIR/../dataset \ + --pretrained_model_dir $HYDRAGNN_ROOT/examples/opf/pretrained_models \ + --pretrained_model_name $PRETRAINED_MODEL \ + --finetune_regime $FT_REGIME \ + ${NO_PRETRAINED:+--no_pretrained} \ + ${MAX_TRAIN_SAMPLES:+--max_train_samples $MAX_TRAIN_SAMPLES} + echo "--- Phase 2 complete ---" +fi + +echo "" +echo "Fine-tuning job finished. Logs: $HYDRAGNN_ROOT/logs/$LOG_NAME/" diff --git a/examples/opf/finetune/plot_finetune_paper_figures.py b/examples/opf/finetune/plot_finetune_paper_figures.py new file mode 100644 index 000000000..f58e4225c --- /dev/null +++ b/examples/opf/finetune/plot_finetune_paper_figures.py @@ -0,0 +1,554 @@ +""" +Generate all fine-tuning paper figures. + +Figures produced (saved to figures/results/ in the LaTeX project): + fig09_ft1_sample_efficiency.pdf FT1 accuracy/F1/AUC vs n (2-arch × 3-metric) + fig10_ft1_convergence.pdf FT1 val-loss vs epoch for HeteroSAGE at 4 n values + fig11_ft3_sample_efficiency.pdf FT3 Va_r2 / Vm_r2 vs n (2-arch × 2-metric) + fig12_ft3_convergence.pdf FT3 val-loss vs epoch for HeteroSAGE at 4 n values + +Usage (from examples/opf/finetune/): + python3.11 plot_finetune_paper_figures.py +""" + +import os, re, json, glob, warnings +import numpy as np +import pandas as pd +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.ticker as ticker + +# ── output directory ───────────────────────────────────────────────────────── +LATEX_DIR = "/lustre/orion/lrn078/proj-shared/hydragnn_opf_tsg_project_with_layer_references" +OUT_DIR = os.path.join(LATEX_DIR, "figures", "results") +os.makedirs(OUT_DIR, exist_ok=True) + +LOGS_DIR = "logs" + +# ── style ──────────────────────────────────────────────────────────────────── +REGIME_ORDER = ["head_only", "partial", "full", "scratch"] +REGIME_LABELS = { + "head_only": "Head-only FT", + "partial": "Partial FT", + "full": "Full FT", + "scratch": "Scratch", +} +REGIME_COLORS = { + "head_only": "#4E79A7", # blue + "partial": "#F28E2B", # orange + "full": "#59A14F", # green + "scratch": "#E15759", # red +} +REGIME_MARKERS = { + "head_only": "o", + "partial": "s", + "full": "^", + "scratch": "D", +} +REGIME_DASHES = { + "head_only": (None, None), # solid + "partial": (4, 2), + "full": (None, None), + "scratch": (2, 2), +} +ARCH_ORDER = ["HeteroSAGE", "HeteroHEAT"] +ARCH_LABELS = {"HeteroSAGE": "HeteroSAGE", "HeteroHEAT": "HeteroHEAT"} + +plt.rcParams.update({ + "figure.dpi": 150, + "font.family": "serif", + "font.size": 9, + "axes.labelsize": 9, + "axes.titlesize": 9, + "legend.fontsize": 8, + "xtick.labelsize": 8, + "ytick.labelsize": 8, + "lines.linewidth": 1.6, + "lines.markersize": 5, + "savefig.bbox": "tight", + "savefig.dpi": 300, +}) + +# ── helpers ─────────────────────────────────────────────────────────────────── + +def _parse_ft1_dirname(name): + """Return (arch, regime, n) or None.""" + m = re.fullmatch( + r"FT1_feasibility_(HeteroSAGE|HeteroHEAT)_" + r"(full|partial|head_only|full_scratch)_n(\d+)", name) + if not m: + return None + arch, regime_raw, n = m.group(1), m.group(2), int(m.group(3)) + regime = "scratch" if regime_raw == "full_scratch" else regime_raw + return arch, regime, n + + +def _parse_ft3_dirname(name): + """Return (arch, regime, n) or None.""" + m = re.fullmatch( + r"finetune_FT3_contingency_(HeteroSAGE|HeteroHEAT)_" + r"(full|partial|head_only|full_scratch)_n(\d+)", name) + if not m: + return None + arch, regime_raw, n = m.group(1), m.group(2), int(m.group(3)) + regime = "scratch" if regime_raw == "full_scratch" else regime_raw + return arch, regime, n + + +def _load_results(log_dir): + """Load results.json; return (meta, test_metrics) or (None, None).""" + path = os.path.join(log_dir, "results.json") + if not os.path.exists(path): + return None, None + with open(path) as f: + d = json.load(f) + return d.get("meta", {}), d.get("test_metrics", {}) + + +def _load_curve(log_dir, tag="validate error"): + """Load training_curve.csv; return (epochs, values) arrays or (None, None).""" + path = os.path.join(log_dir, "training_curve.csv") + if not os.path.exists(path): + return None, None + df = pd.read_csv(path) + sub = df[df["tag"] == tag].copy() + if sub.empty: + return None, None + sub = sub.sort_values("step") + return sub["step"].values, sub["value"].values + + +def _collect_ft1(): + """Collect all FT1 final metrics. Returns DataFrame.""" + rows = [] + for name in sorted(os.listdir(LOGS_DIR)): + parsed = _parse_ft1_dirname(name) + if parsed is None: + continue + arch, regime, n = parsed + log_dir = os.path.join(LOGS_DIR, name) + meta, tm = _load_results(log_dir) + if tm is None: + continue + rows.append({ + "arch": arch, "regime": regime, "n": n, + "accuracy": tm.get("accuracy", np.nan), + "f1": tm.get("f1", np.nan), + "auc_roc": tm.get("auc_roc", np.nan), + "bce": tm.get("bce", np.nan), + "log_dir": log_dir, + }) + return pd.DataFrame(rows) + + +def _collect_ft3(): + """Collect all FT3 final metrics. Returns DataFrame.""" + rows = [] + for name in sorted(os.listdir(LOGS_DIR)): + parsed = _parse_ft3_dirname(name) + if parsed is None: + continue + arch, regime, n = parsed + log_dir = os.path.join(LOGS_DIR, name) + meta, tm = _load_results(log_dir) + if tm is None: + continue + rows.append({ + "arch": arch, "regime": regime, "n": n, + "Va_mse": tm.get("Va_mse", np.nan), + "Va_mae": tm.get("Va_mae", np.nan), + "Va_r2": tm.get("Va_r2", np.nan), + "Vm_mse": tm.get("Vm_mse", np.nan), + "Vm_mae": tm.get("Vm_mae", np.nan), + "Vm_r2": tm.get("Vm_r2", np.nan), + "log_dir": log_dir, + }) + return pd.DataFrame(rows) + + +def _regime_line(ax, df, metric, n_col="n", regimes=None, ylog=False): + """Plot one line per regime on ax.""" + if regimes is None: + regimes = REGIME_ORDER + for regime in regimes: + sub = df[df["regime"] == regime].sort_values(n_col) + if sub.empty: + continue + vals = sub[metric].values + ns = sub[n_col].values + mask = ~np.isnan(vals) + if mask.sum() == 0: + continue + dashes = REGIME_DASHES[regime] + ls = "--" if dashes[0] else "-" + ax.plot(ns[mask], vals[mask], + marker=REGIME_MARKERS[regime], + color=REGIME_COLORS[regime], + linestyle=ls, + label=REGIME_LABELS[regime], + zorder=3) + ax.set_xscale("log") + if ylog: + ax.set_yscale("log") + ax.xaxis.set_major_formatter(ticker.FuncFormatter( + lambda x, _: f"{int(x):,}")) + ax.tick_params(axis="x", rotation=30) + ax.grid(True, which="major", linestyle=":", alpha=0.5) + ax.grid(True, which="minor", linestyle=":", alpha=0.2) + + +# ── Fig 09 — FT1 sample efficiency ─────────────────────────────────────────── + +def plot_ft1_sample_efficiency(df): + metrics = [ + ("accuracy", "Accuracy", False, (0, 1.05)), + ("f1", "F1 score", False, (0, 1.05)), + ("auc_roc", "AUC-ROC", False, (0, 1.05)), + ] + fig, axes = plt.subplots(2, 3, figsize=(7.2, 4.5), sharex=False) + fig.subplots_adjust(hspace=0.45, wspace=0.35) + + for r_idx, arch in enumerate(ARCH_ORDER): + sub_arch = df[df["arch"] == arch] + for c_idx, (met, label, ylog, ylim) in enumerate(metrics): + ax = axes[r_idx, c_idx] + _regime_line(ax, sub_arch, met, ylog=ylog) + ax.set_ylim(ylim) + if r_idx == 1: + ax.set_xlabel("Training samples $n$") + if c_idx == 0: + ax.set_ylabel(label) + else: + ax.set_ylabel(label) + ax.set_title(f"{ARCH_LABELS[arch]}", fontweight="bold") + if r_idx == 0 and c_idx == 2: + ax.legend(loc="lower right", framealpha=0.85) + + # column titles + for c_idx, (_, label, _, _) in enumerate(metrics): + axes[0, c_idx].set_title( + f"{label}\n{ARCH_LABELS[ARCH_ORDER[0]]}", fontweight="bold") + axes[1, c_idx].set_title( + f"{ARCH_LABELS[ARCH_ORDER[1]]}", fontweight="bold") + + fig.suptitle( + "FT1 — Feasibility Classification: Sample Efficiency", + fontsize=10, fontweight="bold", y=1.01) + + path = os.path.join(OUT_DIR, "fig09_ft1_sample_efficiency.pdf") + fig.savefig(path) + fig.savefig(path.replace(".pdf", ".png")) + plt.close(fig) + print(f"Saved {path}") + + +# ── Fig 10 — FT1 convergence (HeteroSAGE) ──────────────────────────────────── + +def plot_ft1_convergence(df_ft1): + arch = "HeteroSAGE" + n_vals = [100, 1000, 10000, 50000] + n_labels = {100: "$n=100$", 1000: "$n=1{,}000$", + 10000: "$n=10{,}000$", 50000: "$n=50{,}000$"} + + fig, axes = plt.subplots(1, 4, figsize=(7.2, 2.4), sharey=False) + fig.subplots_adjust(wspace=0.38) + + for col, n in enumerate(n_vals): + ax = axes[col] + sub = df_ft1[(df_ft1["arch"] == arch) & (df_ft1["n"] == n)] + plotted = False + for regime in REGIME_ORDER: + row = sub[sub["regime"] == regime] + if row.empty: + continue + log_dir = row.iloc[0]["log_dir"] + epochs, vals = _load_curve(log_dir, tag="validate error") + if epochs is None: + continue + dashes = REGIME_DASHES[regime] + ls = "--" if dashes[0] else "-" + ax.plot(epochs + 1, vals, + color=REGIME_COLORS[regime], + linestyle=ls, + label=REGIME_LABELS[regime], + zorder=3) + plotted = True + + ax.set_title(n_labels[n], fontsize=9) + ax.set_xlabel("Epoch") + if col == 0: + ax.set_ylabel("Val. BCE loss") + ax.set_yscale("log") + ax.grid(True, which="major", linestyle=":", alpha=0.5) + ax.grid(True, which="minor", linestyle=":", alpha=0.2) + if col == 3 and plotted: + ax.legend(loc="upper right", fontsize=7.5, framealpha=0.85) + + fig.suptitle( + "FT1 — Training Convergence (HeteroSAGE, validation loss)", + fontsize=10, fontweight="bold", y=1.03) + + path = os.path.join(OUT_DIR, "fig10_ft1_convergence.pdf") + fig.savefig(path) + fig.savefig(path.replace(".pdf", ".png")) + plt.close(fig) + print(f"Saved {path}") + + +# ── Fig 11 — FT3 sample efficiency ─────────────────────────────────────────── + +def plot_ft3_sample_efficiency(df): + metrics = [ + ("Va_r2", r"$R^2$ (voltage angle $V_a$)", False, (-0.1, 1.05)), + ("Vm_r2", r"$R^2$ (voltage magnitude $V_m$)", False, (-0.1, 1.05)), + ] + fig, axes = plt.subplots(2, 2, figsize=(6.0, 4.5), sharex=False) + fig.subplots_adjust(hspace=0.50, wspace=0.40) + + for r_idx, arch in enumerate(ARCH_ORDER): + sub_arch = df[df["arch"] == arch] + for c_idx, (met, label, ylog, ylim) in enumerate(metrics): + ax = axes[r_idx, c_idx] + _regime_line(ax, sub_arch, met, ylog=ylog) + ax.set_ylim(ylim) + ax.axhline(0, color="gray", linewidth=0.8, linestyle=":", zorder=0) + ax.axhline(1, color="gray", linewidth=0.8, linestyle=":", zorder=0) + if r_idx == 1: + ax.set_xlabel("Training samples $n$") + ax.set_ylabel(label) + title_str = label.split("(")[0].strip() + f"\n{ARCH_LABELS[arch]}" + ax.set_title(title_str, fontweight="bold") + if r_idx == 1 and c_idx == 1: + ax.legend(loc="lower right", framealpha=0.85) + + fig.suptitle( + r"FT3 — N-1 Contingency OPF Regression: Sample Efficiency ($R^2$)", + fontsize=10, fontweight="bold", y=1.01) + + path = os.path.join(OUT_DIR, "fig11_ft3_sample_efficiency.pdf") + fig.savefig(path) + fig.savefig(path.replace(".pdf", ".png")) + plt.close(fig) + print(f"Saved {path}") + + +# ── Fig 11b — FT3 sample efficiency MSE (supplemental) ─────────────────────── + +def plot_ft3_sample_efficiency_mse(df): + metrics = [ + ("Va_mse", r"MSE (voltage angle $V_a$)", True), + ("Vm_mse", r"MSE (voltage magnitude $V_m$)", True), + ] + fig, axes = plt.subplots(2, 2, figsize=(6.0, 4.5), sharex=False) + fig.subplots_adjust(hspace=0.50, wspace=0.40) + + for r_idx, arch in enumerate(ARCH_ORDER): + sub_arch = df[df["arch"] == arch] + for c_idx, (met, label, ylog) in enumerate(metrics): + ax = axes[r_idx, c_idx] + _regime_line(ax, sub_arch, met, ylog=ylog) + if r_idx == 1: + ax.set_xlabel("Training samples $n$") + ax.set_ylabel(label) + title_str = label.split("(")[0].strip() + f"\n{ARCH_LABELS[arch]}" + ax.set_title(title_str, fontweight="bold") + if r_idx == 1 and c_idx == 1: + ax.legend(loc="upper right", framealpha=0.85) + + fig.suptitle( + r"FT3 — N-1 Contingency OPF Regression: Sample Efficiency (MSE)", + fontsize=10, fontweight="bold", y=1.01) + + path = os.path.join(OUT_DIR, "fig11b_ft3_sample_efficiency_mse.pdf") + fig.savefig(path) + fig.savefig(path.replace(".pdf", ".png")) + plt.close(fig) + print(f"Saved {path}") + + +# ── Fig 12 — FT3 convergence (HeteroSAGE) ──────────────────────────────────── + +def plot_ft3_convergence(df_ft3): + arch = "HeteroSAGE" + n_vals = [100, 1000, 10000, 50000] + n_labels = {100: "$n=100$", 1000: "$n=1{,}000$", + 10000: "$n=10{,}000$", 50000: "$n=50{,}000$"} + + fig, axes = plt.subplots(1, 4, figsize=(7.2, 2.4), sharey=False) + fig.subplots_adjust(wspace=0.38) + + for col, n in enumerate(n_vals): + ax = axes[col] + sub = df_ft3[(df_ft3["arch"] == arch) & (df_ft3["n"] == n)] + plotted = False + for regime in REGIME_ORDER: + row = sub[sub["regime"] == regime] + if row.empty: + continue + log_dir = row.iloc[0]["log_dir"] + epochs, vals = _load_curve(log_dir, tag="validate error") + if epochs is None: + continue + dashes = REGIME_DASHES[regime] + ls = "--" if dashes[0] else "-" + ax.plot(epochs + 1, vals, + color=REGIME_COLORS[regime], + linestyle=ls, + label=REGIME_LABELS[regime], + zorder=3) + plotted = True + + ax.set_title(n_labels[n], fontsize=9) + ax.set_xlabel("Epoch") + if col == 0: + ax.set_ylabel("Val. MSE loss") + ax.set_yscale("log") + ax.grid(True, which="major", linestyle=":", alpha=0.5) + ax.grid(True, which="minor", linestyle=":", alpha=0.2) + if col == 3 and plotted: + ax.legend(loc="upper right", fontsize=7.5, framealpha=0.85) + + fig.suptitle( + "FT3 — Training Convergence (HeteroSAGE, validation loss)", + fontsize=10, fontweight="bold", y=1.03) + + path = os.path.join(OUT_DIR, "fig12_ft3_convergence.pdf") + fig.savefig(path) + fig.savefig(path.replace(".pdf", ".png")) + plt.close(fig) + print(f"Saved {path}") + + +# ── Fig 13 — FT3 convergence HeteroHEAT (side-by-side with SAGE) ───────────── + +def plot_ft3_convergence_both_archs(df_ft3): + """2-row × 4-col: top=HeteroSAGE, bottom=HeteroHEAT; cols=n values.""" + n_vals = [100, 1000, 10000, 50000] + n_labels = {100: "$n=100$", 1000: "$n=1{,}000$", + 10000: "$n=10{,}000$", 50000: "$n=50{,}000$"} + + fig, axes = plt.subplots(2, 4, figsize=(7.2, 4.5), sharey=False) + fig.subplots_adjust(hspace=0.50, wspace=0.38) + + for r_idx, arch in enumerate(ARCH_ORDER): + for col, n in enumerate(n_vals): + ax = axes[r_idx, col] + sub = df_ft3[(df_ft3["arch"] == arch) & (df_ft3["n"] == n)] + plotted = False + for regime in REGIME_ORDER: + row = sub[sub["regime"] == regime] + if row.empty: + continue + log_dir = row.iloc[0]["log_dir"] + epochs, vals = _load_curve(log_dir, tag="validate error") + if epochs is None: + continue + dashes = REGIME_DASHES[regime] + ls = "--" if dashes[0] else "-" + ax.plot(epochs + 1, vals, + color=REGIME_COLORS[regime], + linestyle=ls, + label=REGIME_LABELS[regime], + zorder=3) + plotted = True + + if r_idx == 0: + ax.set_title(n_labels[n], fontsize=9) + ax.set_xlabel("Epoch") + if col == 0: + ax.set_ylabel(f"{ARCH_LABELS[arch]}\nVal. MSE loss") + ax.set_yscale("log") + ax.grid(True, which="major", linestyle=":", alpha=0.5) + ax.grid(True, which="minor", linestyle=":", alpha=0.2) + if r_idx == 1 and col == 3 and plotted: + ax.legend(loc="upper right", fontsize=7.5, framealpha=0.85) + + fig.suptitle( + "FT3 — Training Convergence by Architecture and Dataset Size", + fontsize=10, fontweight="bold", y=1.01) + + path = os.path.join(OUT_DIR, "fig13_ft3_convergence_both_archs.pdf") + fig.savefig(path) + fig.savefig(path.replace(".pdf", ".png")) + plt.close(fig) + print(f"Saved {path}") + + +# ── Fig 14 — FT1 convergence both archs ────────────────────────────────────── + +def plot_ft1_convergence_both_archs(df_ft1): + n_vals = [100, 1000, 10000, 50000] + n_labels = {100: "$n=100$", 1000: "$n=1{,}000$", + 10000: "$n=10{,}000$", 50000: "$n=50{,}000$"} + + fig, axes = plt.subplots(2, 4, figsize=(7.2, 4.5), sharey=False) + fig.subplots_adjust(hspace=0.50, wspace=0.38) + + for r_idx, arch in enumerate(ARCH_ORDER): + for col, n in enumerate(n_vals): + ax = axes[r_idx, col] + sub = df_ft1[(df_ft1["arch"] == arch) & (df_ft1["n"] == n)] + plotted = False + for regime in REGIME_ORDER: + row = sub[sub["regime"] == regime] + if row.empty: + continue + log_dir = row.iloc[0]["log_dir"] + epochs, vals = _load_curve(log_dir, tag="validate error") + if epochs is None: + continue + dashes = REGIME_DASHES[regime] + ls = "--" if dashes[0] else "-" + ax.plot(epochs + 1, vals, + color=REGIME_COLORS[regime], + linestyle=ls, + label=REGIME_LABELS[regime], + zorder=3) + plotted = True + + if r_idx == 0: + ax.set_title(n_labels[n], fontsize=9) + ax.set_xlabel("Epoch") + if col == 0: + ax.set_ylabel(f"{ARCH_LABELS[arch]}\nVal. BCE loss") + ax.set_yscale("log") + ax.grid(True, which="major", linestyle=":", alpha=0.5) + ax.grid(True, which="minor", linestyle=":", alpha=0.2) + if r_idx == 1 and col == 3 and plotted: + ax.legend(loc="upper right", fontsize=7.5, framealpha=0.85) + + fig.suptitle( + "FT1 — Training Convergence by Architecture and Dataset Size", + fontsize=10, fontweight="bold", y=1.01) + + path = os.path.join(OUT_DIR, "fig14_ft1_convergence_both_archs.pdf") + fig.savefig(path) + fig.savefig(path.replace(".pdf", ".png")) + plt.close(fig) + print(f"Saved {path}") + + +# ── main ───────────────────────────────────────────────────────────────────── + +if __name__ == "__main__": + print("Loading FT1 results …") + df_ft1 = _collect_ft1() + print(f" {len(df_ft1)} runs found") + print(df_ft1.groupby(["arch", "regime"]).size().to_string()) + + print("\nLoading FT3 results …") + df_ft3 = _collect_ft3() + print(f" {len(df_ft3)} runs found") + print(df_ft3.groupby(["arch", "regime"]).size().to_string()) + + print("\nGenerating figures …") + plot_ft1_sample_efficiency(df_ft1) + plot_ft1_convergence(df_ft1) + plot_ft1_convergence_both_archs(df_ft1) + plot_ft3_sample_efficiency(df_ft3) + plot_ft3_sample_efficiency_mse(df_ft3) + plot_ft3_convergence(df_ft3) + plot_ft3_convergence_both_archs(df_ft3) + + print("\nAll figures written to:", OUT_DIR) diff --git a/examples/opf/finetune/plot_ft1_bce_epoch_curves.py b/examples/opf/finetune/plot_ft1_bce_epoch_curves.py new file mode 100644 index 000000000..4e2b30aa6 --- /dev/null +++ b/examples/opf/finetune/plot_ft1_bce_epoch_curves.py @@ -0,0 +1,218 @@ +#!/usr/bin/env python3 +"""Plot FT1 BCE loss curves vs epoch for HeteroSAGE and HeteroHEAT. + +This script scans FT1 run directories under logs/, parses per-epoch losses from +run.log, and aggregates runs by fine-tuning strategy (full, partial, head_only, +scratch). It generates one figure per architecture with mean curve and +/- 1 std +band across all N-regimes. + +Outputs are written as both PNG and PDF for manuscript use. +""" + +import argparse +import re +from collections import defaultdict +from pathlib import Path +from statistics import mean +from typing import Dict, List, Tuple + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + + +FT1_RE = re.compile( + r"^FT1_feasibility_(?P[A-Za-z0-9]+)_(?Pfull|partial|head_only)" + r"(?P_scratch)?_n(?P\d+)$" +) + +EPOCH_RE = re.compile( + r"Epoch:\s*(?P\d+),\s*" + r"Train Loss:\s*(?P[-+0-9.eE]+),\s*" + r"Val Loss:\s*(?P[-+0-9.eE]+),\s*" + r"Test Loss:\s*(?P[-+0-9.eE]+)" +) + +METHOD_ORDER = ["full", "partial", "head_only", "scratch"] +METHOD_LABEL = { + "full": "Full", + "partial": "Partial", + "head_only": "Head-only", + "scratch": "Scratch", +} +METHOD_STYLE = { + "full": {"color": "#1f77b4"}, + "partial": {"color": "#ff7f0e"}, + "head_only": {"color": "#2ca02c"}, + "scratch": {"color": "#d62728"}, +} + + +def parse_run_log(run_log: Path, split: str) -> Dict[int, float]: + """Return epoch->loss parsed from run.log for the requested split.""" + loss_by_epoch = {} + with run_log.open("r", encoding="utf-8", errors="ignore") as fh: + for line in fh: + m = EPOCH_RE.search(line) + if not m: + continue + epoch = int(m.group("epoch")) + loss = float(m.group(split)) + loss_by_epoch[epoch] = loss + return loss_by_epoch + + +def collect(logs_dir: Path, split: str) -> Dict[str, Dict[str, List[Dict[int, float]]]]: + """Collect per-run epoch curves organized as arch -> method -> list(curves).""" + curves = defaultdict( + lambda: defaultdict(list) + ) + + for d in sorted(logs_dir.iterdir()): + if not d.is_dir(): + continue + m = FT1_RE.match(d.name) + if not m: + continue + + arch = m.group("arch") + method = "scratch" if m.group("scratch") else m.group("regime") + + run_log = d / "run.log" + if not run_log.is_file(): + continue + + curve = parse_run_log(run_log, split=split) + if curve: + curves[arch][method].append(curve) + + return curves + + +def summarize_curves(curves: List[Dict[int, float]]) -> Tuple[List[int], List[float], List[float]]: + """Return epochs, mean_loss, std_loss from a list of epoch->loss curves.""" + by_epoch = defaultdict(list) + for c in curves: + for e, v in c.items(): + by_epoch[e].append(v) + + epochs = sorted(by_epoch) + means = [mean(by_epoch[e]) for e in epochs] + stds = [] + for e in epochs: + vals = by_epoch[e] + mu = mean(vals) + var = sum((x - mu) ** 2 for x in vals) / len(vals) + stds.append(var ** 0.5) + return epochs, means, stds + + +def configure_pub_style() -> None: + """Set plotting defaults suitable for a two-column IEEE manuscript.""" + plt.rcParams.update( + { + "font.family": "DejaVu Sans", + "font.size": 16, + "axes.labelsize": 20, + "axes.titlesize": 22, + "axes.titleweight": "bold", + "xtick.labelsize": 16, + "ytick.labelsize": 16, + "legend.fontsize": 15, + "legend.title_fontsize": 15, + "lines.linewidth": 3.0, + } + ) + + +def plot_architecture( + arch: str, + method_curves: Dict[str, List[Dict[int, float]]], + split_label: str, + out_base: Path, + dpi: int, +) -> None: + """Plot one architecture figure and save PNG+PDF.""" + fig, ax = plt.subplots(figsize=(7.2, 4.8), constrained_layout=True) + + plotted_any = False + for method in METHOD_ORDER: + curves = method_curves.get(method, []) + if not curves: + continue + + epochs, means, stds = summarize_curves(curves) + color = METHOD_STYLE[method]["color"] + label = f"{METHOD_LABEL[method]} (n={len(curves)})" + + ax.plot(epochs, means, color=color, label=label) + lower = [max(m - s, 1e-12) for m, s in zip(means, stds)] + upper = [m + s for m, s in zip(means, stds)] + ax.fill_between(epochs, lower, upper, color=color, alpha=0.18) + plotted_any = True + + if not plotted_any: + plt.close(fig) + return + + ax.set_yscale("log") + ax.set_xlabel("Epoch") + ax.set_ylabel(f"{split_label} BCE Loss") + ax.set_title(f"FT1 {arch}: {split_label} BCE vs Epoch") + ax.grid(True, which="both", linestyle=":", linewidth=0.9, alpha=0.8) + ax.legend(title="Fine-tuning Strategy", loc="best", frameon=True) + + png_path = out_base.with_suffix(".png") + pdf_path = out_base.with_suffix(".pdf") + fig.savefig(png_path, dpi=dpi) + fig.savefig(pdf_path) + plt.close(fig) + print(f"[plot] {png_path}") + print(f"[plot] {pdf_path}") + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--logs", + type=Path, + default=Path(__file__).parent / "logs", + help="Path to FT run logs directory", + ) + ap.add_argument( + "--out", + type=Path, + default=Path(__file__).parent / "plots", + help="Directory to save figure files", + ) + ap.add_argument( + "--split", + choices=["train", "val", "test"], + default="val", + help="Which BCE split from run.log to plot", + ) + ap.add_argument( + "--dpi", + type=int, + default=400, + help="PNG export DPI (use >=300 for publication)", + ) + args = ap.parse_args() + + configure_pub_style() + args.out.mkdir(parents=True, exist_ok=True) + + curves = collect(args.logs, split=args.split) + split_label = {"train": "Train", "val": "Validation", "test": "Test"}[args.split] + + for arch in ["HeteroSAGE", "HeteroHEAT"]: + if arch not in curves: + print(f"[warn] no FT1 curves found for {arch}") + continue + out_base = args.out / f"FT1_{arch}_{args.split}_bce_vs_epoch" + plot_architecture(arch, curves[arch], split_label, out_base, args.dpi) + + +if __name__ == "__main__": + main() diff --git a/examples/opf/finetune/plot_ft_results.py b/examples/opf/finetune/plot_ft_results.py new file mode 100644 index 000000000..b9cc4a365 --- /dev/null +++ b/examples/opf/finetune/plot_ft_results.py @@ -0,0 +1,440 @@ +"""Generate comparison plots for FT1 and FT3 fine-tuning experiments. + +Reads the summary JSON produced by collect_results.py and writes a set of +publication-quality figures to results/figures/. + +Figures produced +---------------- +FT1 (feasibility classification): + ft1_metrics_bar.pdf Grouped bar chart: Accuracy, F1, AUC-ROC by regime × arch + ft1_roc_curves.pdf ROC curves for every run on the same axes + ft1_learning_curves.pdf Train/val loss curves per run (grid layout) + +FT3 (N-1 contingency regression): + ft3_mse_bar.pdf Grouped bar chart: Va MSE, Vm MSE by regime × arch + ft3_r2_bar.pdf Grouped bar chart: Va R², Vm R² by regime × arch + ft3_learning_curves.pdf Train/val loss curves per run (grid layout) + ft3_scatter_best.pdf Pred vs actual scatter for best model (Va and Vm) + +Usage:: + + # From examples/opf/finetune/ + python plot_ft_results.py + python plot_ft_results.py --summary results/ft1_ft3_summary.json \\ + --out_dir results/figures +""" + +import argparse +import json +import os +from itertools import groupby + +import matplotlib +matplotlib.use("Agg") # non-interactive backend for headless runs +import matplotlib.pyplot as plt +import matplotlib.patches as mpatches +import numpy as np + + +# ───────────────────────────────────────────────────────────────────────────── +# Style +# ───────────────────────────────────────────────────────────────────────────── + +# Colour palette — one colour per regime + baseline +REGIME_COLORS = { + "head_only": "#4E79A7", + "partial": "#F28E2B", + "full": "#59A14F", + "full_baseline": "#E15759", # full + no_pretrained + "baseline": "#E15759", +} +ARCH_HATCHES = {"HeteroSAGE": "", "HeteroHEAT": "//"} + +LABEL_MAP = { + "head_only": "Head-only (FT)", + "partial": "Partial (FT)", + "full": "Full (FT)", + "baseline": "Scratch (baseline)", + "full_baseline": "Scratch (baseline)", +} + +plt.rcParams.update({ + "figure.dpi": 150, + "font.size": 11, + "axes.labelsize": 12, + "axes.titlesize": 13, + "legend.fontsize": 10, + "xtick.labelsize": 10, + "ytick.labelsize": 10, + "savefig.bbox": "tight", +}) + + +# ───────────────────────────────────────────────────────────────────────────── +# Data helpers +# ───────────────────────────────────────────────────────────────────────────── + +def _load_summary(path: str) -> list[dict]: + with open(path) as fh: + return json.load(fh) + + +def _split(runs: list[dict], strategy_prefix: str) -> list[dict]: + return [r for r in runs if r["meta"].get("ft_strategy", "").startswith(strategy_prefix)] + + +def _label(run: dict) -> str: + meta = run["meta"] + arch = meta.get("arch", "?") + regime = meta.get("regime", "full") + pretrained = meta.get("pretrained", True) + if not pretrained: + return f"{arch}_baseline" + return f"{arch}_{regime}" + + +def _regime_key(run: dict) -> str: + meta = run["meta"] + pretrained = meta.get("pretrained", True) + regime = meta.get("regime", "full") + return regime if pretrained else "baseline" + + +def _color(run: dict) -> str: + return REGIME_COLORS.get(_regime_key(run), "#888888") + + +def _hatch(run: dict) -> str: + arch = run["meta"].get("arch", "") + return ARCH_HATCHES.get(arch, "") + + +# ───────────────────────────────────────────────────────────────────────────── +# FT1 plots +# ───────────────────────────────────────────────────────────────────────────── + +def plot_ft1_metrics_bar(runs: list[dict], out_dir: str) -> None: + """Grouped bar chart comparing Accuracy, F1, AUC-ROC for every FT1 run.""" + metrics_labels = [("accuracy", "Accuracy"), ("f1", "F1 Score"), ("auc_roc", "AUC-ROC")] + n_metrics = len(metrics_labels) + n_runs = len(runs) + + fig, axes = plt.subplots(1, n_metrics, figsize=(5 * n_metrics, 5), sharey=False) + if n_metrics == 1: + axes = [axes] + + for ax, (metric_key, metric_title) in zip(axes, metrics_labels): + xs = np.arange(n_runs) + heights = [ + r["test_metrics"].get(metric_key) or 0.0 + for r in runs + ] + bar_labels = [_label(r) for r in runs] + colors = [_color(r) for r in runs] + hatches = [_hatch(r) for r in runs] + + bars = ax.bar(xs, heights, color=colors, edgecolor="black", linewidth=0.7) + for bar, h_pat in zip(bars, hatches): + bar.set_hatch(h_pat) + + ax.set_xticks(xs) + ax.set_xticklabels(bar_labels, rotation=30, ha="right") + ax.set_title(metric_title) + ax.set_ylim(0, 1.05) + ax.axhline(1.0, color="black", linewidth=0.5, linestyle="--") + + fig.suptitle("FT1 – Feasibility Classification: Test Metrics", fontsize=14, y=1.02) + + # Legend for colours (regime) and hatches (arch) + legend_patches = [ + mpatches.Patch(facecolor=v, edgecolor="black", label=LABEL_MAP.get(k, k)) + for k, v in REGIME_COLORS.items() + ] + legend_patches += [ + mpatches.Patch(facecolor="white", hatch=h, edgecolor="black", + label=arch) + for arch, h in ARCH_HATCHES.items() + ] + axes[-1].legend(handles=legend_patches, bbox_to_anchor=(1.05, 1), loc="upper left", + borderaxespad=0, frameon=True) + + _save(fig, out_dir, "ft1_metrics_bar") + + +def plot_ft1_roc_curves(runs: list[dict], out_dir: str) -> None: + """ROC curves for all FT1 runs on the same axes.""" + try: + from sklearn.metrics import roc_curve, auc + except ImportError: + print("[plot] sklearn not available — skipping ROC plot.") + return + + fig, ax = plt.subplots(figsize=(6, 5)) + ax.plot([0, 1], [0, 1], "k--", linewidth=0.8, label="Random") + + for run in runs: + probs = run.get("probs", []) + labels = run.get("labels", []) + if not probs or len(set(labels)) < 2: + continue + fpr, tpr, _ = roc_curve(labels, probs) + roc_auc = auc(fpr, tpr) + lbl = _label(run) + ax.plot(fpr, tpr, label=f"{lbl} (AUC={roc_auc:.3f})", + color=_color(run), + linestyle="--" if "HeteroHEAT" in lbl else "-", + linewidth=1.8) + + ax.set_xlabel("False Positive Rate") + ax.set_ylabel("True Positive Rate") + ax.set_title("FT1 – Feasibility Classification: ROC Curves") + ax.legend(loc="lower right", fontsize=9) + ax.grid(alpha=0.3) + + _save(fig, out_dir, "ft1_roc_curves") + + +def plot_ft1_learning_curves(runs: list[dict], out_dir: str) -> None: + _plot_learning_curves(runs, out_dir, prefix="ft1", + title="FT1 – Feasibility Classification: Training Curves") + + +# ───────────────────────────────────────────────────────────────────────────── +# FT3 plots +# ───────────────────────────────────────────────────────────────────────────── + +def plot_ft3_mse_bar(runs: list[dict], out_dir: str) -> None: + """Grouped bar chart: Va MSE and Vm MSE for every FT3 run.""" + fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=False) + _metric_bar(axes[0], runs, "Va_mse", "Va Voltage Angle MSE") + _metric_bar(axes[1], runs, "Vm_mse", "Vm Voltage Magnitude MSE") + fig.suptitle("FT3 – N-1 Contingency Regression: MSE Comparison", fontsize=14, y=1.02) + _add_legend(axes[-1]) + _save(fig, out_dir, "ft3_mse_bar") + + +def plot_ft3_r2_bar(runs: list[dict], out_dir: str) -> None: + """Grouped bar chart: Va R² and Vm R² for every FT3 run.""" + fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=False) + _metric_bar(axes[0], runs, "Va_r2", "Va Voltage Angle R²", lower_better=False) + _metric_bar(axes[1], runs, "Vm_r2", "Vm Voltage Magnitude R²", lower_better=False) + fig.suptitle("FT3 – N-1 Contingency Regression: R² Comparison", fontsize=14, y=1.02) + _add_legend(axes[-1]) + _save(fig, out_dir, "ft3_r2_bar") + + +def plot_ft3_learning_curves(runs: list[dict], out_dir: str) -> None: + _plot_learning_curves(runs, out_dir, prefix="ft3", + title="FT3 – N-1 Contingency Regression: Training Curves") + + +def plot_ft3_scatter_best(runs: list[dict], out_dir: str) -> None: + """Pred vs actual scatter for the best (lowest Va_mse) FT3 run.""" + valid = [r for r in runs if r["test_metrics"].get("Va_mse") is not None + and r.get("preds_sample")] + if not valid: + print("[plot] No FT3 runs with scatter data — skipping scatter plot.") + return + + best = min(valid, key=lambda r: r["test_metrics"]["Va_mse"]) + preds = np.array(best["preds_sample"]) # [N, 2] + targets = np.array(best["targets_sample"]) + + if preds.ndim == 1 or preds.shape[1] < 2: + print("[plot] Unexpected preds_sample shape — skipping scatter.") + return + + fig, axes = plt.subplots(1, 2, figsize=(10, 4)) + for ax, dim, name in zip(axes, [0, 1], ["Va (voltage angle)", "Vm (voltage magnitude)"]): + p, t = preds[:, dim], targets[:, dim] + ax.scatter(t, p, alpha=0.3, s=8, color="#4E79A7") + lim = [min(t.min(), p.min()), max(t.max(), p.max())] + ax.plot(lim, lim, "r--", linewidth=1, label="y = x") + ax.set_xlabel(f"Actual {name}") + ax.set_ylabel(f"Predicted {name}") + ax.set_title(f"{name}") + ax.legend(fontsize=9) + + fig.suptitle(f"FT3 – Best model ({_label(best)}): Pred vs Actual", + fontsize=13, y=1.02) + _save(fig, out_dir, "ft3_scatter_best") + + +# ───────────────────────────────────────────────────────────────────────────── +# Shared helpers +# ───────────────────────────────────────────────────────────────────────────── + +def _metric_bar(ax, runs: list[dict], key: str, title: str, + lower_better: bool = True) -> None: + n_runs = len(runs) + xs = np.arange(n_runs) + heights = [r["test_metrics"].get(key) for r in runs] + bar_labels = [_label(r) for r in runs] + colors = [_color(r) for r in runs] + hatches = [_hatch(r) for r in runs] + + bars = ax.bar(xs, [h or 0.0 for h in heights], + color=colors, edgecolor="black", linewidth=0.7) + for bar, h_pat in zip(bars, hatches): + bar.set_hatch(h_pat) + + # Annotate missing values + for x, h in zip(xs, heights): + if h is None: + ax.text(x, 0.01, "N/A", ha="center", va="bottom", fontsize=8, color="gray") + + ax.set_xticks(xs) + ax.set_xticklabels(bar_labels, rotation=30, ha="right") + ax.set_title(title) + if lower_better: + ax.set_ylabel("MSE ↓") + else: + ax.set_ylabel("R² ↑") + ax.axhline(1.0, color="black", linewidth=0.5, linestyle="--") + + +def _add_legend(ax) -> None: + patches = [ + mpatches.Patch(facecolor=v, edgecolor="black", label=LABEL_MAP.get(k, k)) + for k, v in REGIME_COLORS.items() + ] + patches += [ + mpatches.Patch(facecolor="white", hatch=h, edgecolor="black", label=arch) + for arch, h in ARCH_HATCHES.items() + ] + ax.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc="upper left", + borderaxespad=0, frameon=True) + + +def _plot_learning_curves(runs: list[dict], out_dir: str, + prefix: str, title: str) -> None: + n = len(runs) + if n == 0: + return + ncols = min(4, n) + nrows = (n + ncols - 1) // ncols + fig, axes = plt.subplots(nrows, ncols, figsize=(4.5 * ncols, 3.5 * nrows), + squeeze=False) + + for idx, run in enumerate(runs): + ax = axes[idx // ncols][idx % ncols] + curve = run.get("training_curve", {}) + + train_vals = curve.get("train error", []) + val_vals = curve.get("validate error", []) + + if train_vals: + epochs = [ep for ep, _ in train_vals] + vals = [v for _, v in train_vals] + ax.plot(epochs, vals, label="train", color="#4E79A7", linewidth=1.5) + if val_vals: + epochs = [ep for ep, _ in val_vals] + vals = [v for _, v in val_vals] + ax.plot(epochs, vals, label="val", color="#F28E2B", + linewidth=1.5, linestyle="--") + + ax.set_title(_label(run), fontsize=9) + ax.set_xlabel("Epoch", fontsize=8) + ax.set_ylabel("Loss", fontsize=8) + ax.legend(fontsize=7) + ax.grid(alpha=0.3) + + # Hide unused subplots + for idx in range(n, nrows * ncols): + axes[idx // ncols][idx % ncols].set_visible(False) + + fig.suptitle(title, fontsize=13, y=1.01) + fig.tight_layout() + _save(fig, out_dir, f"{prefix}_learning_curves") + + +def _save(fig, out_dir: str, name: str) -> None: + os.makedirs(out_dir, exist_ok=True) + for ext in ("pdf", "png"): + path = os.path.join(out_dir, f"{name}.{ext}") + fig.savefig(path, bbox_inches="tight") + print(f"[plot] Saved {path}") + plt.close(fig) + + +# ───────────────────────────────────────────────────────────────────────────── +# Combined FT1 vs FT3 comparison +# ───────────────────────────────────────────────────────────────────────────── + +def plot_combined_summary(ft1_runs: list[dict], ft3_runs: list[dict], + out_dir: str) -> None: + """Single-page summary: FT1 F1/AUC + FT3 Va R²/Vm R² — 4 panels.""" + if not ft1_runs and not ft3_runs: + return + + fig, axes = plt.subplots(1, 4, figsize=(22, 5)) + + if ft1_runs: + _metric_bar(axes[0], ft1_runs, "f1", "FT1 — F1 Score") + _metric_bar(axes[1], ft1_runs, "auc_roc", "FT1 — AUC-ROC", + lower_better=False) + if ft3_runs: + _metric_bar(axes[2], ft3_runs, "Va_r2", "FT3 — Va R²", + lower_better=False) + _metric_bar(axes[3], ft3_runs, "Vm_r2", "FT3 — Vm R²", + lower_better=False) + + fig.suptitle("Fine-tuning vs Baseline — FT1 & FT3 Summary", fontsize=14) + _add_legend(axes[-1]) + _save(fig, out_dir, "combined_summary") + + +# ───────────────────────────────────────────────────────────────────────────── +# CLI +# ───────────────────────────────────────────────────────────────────────────── + +def main(): + parser = argparse.ArgumentParser( + description="Generate comparison plots for FT1 and FT3 experiments." + ) + _ft_dir = os.path.dirname(os.path.abspath(__file__)) + _res_dir = os.path.join(_ft_dir, "results") + parser.add_argument( + "--summary", + default=os.path.join(_res_dir, "ft1_ft3_summary.json"), + help="Path to ft1_ft3_summary.json (produced by collect_results.py)", + ) + parser.add_argument( + "--out_dir", + default=os.path.join(_res_dir, "figures"), + help="Output directory for figures (default: results/figures/)", + ) + args = parser.parse_args() + + if not os.path.isfile(args.summary): + print(f"[plot] Summary file not found: {args.summary}") + print(" Run collect_results.py first.") + return + + all_runs = _load_summary(args.summary) + ft1_runs = _split(all_runs, "FT1") + ft3_runs = _split(all_runs, "FT3") + + print(f"[plot] FT1 runs: {len(ft1_runs)} FT3 runs: {len(ft3_runs)}") + + # ── FT1 figures ──────────────────────────────────────────────────────── + if ft1_runs: + plot_ft1_metrics_bar(ft1_runs, args.out_dir) + plot_ft1_roc_curves(ft1_runs, args.out_dir) + plot_ft1_learning_curves(ft1_runs, args.out_dir) + + # ── FT3 figures ──────────────────────────────────────────────────────── + if ft3_runs: + plot_ft3_mse_bar(ft3_runs, args.out_dir) + plot_ft3_r2_bar(ft3_runs, args.out_dir) + plot_ft3_learning_curves(ft3_runs, args.out_dir) + plot_ft3_scatter_best(ft3_runs, args.out_dir) + + # ── Combined summary ─────────────────────────────────────────────────── + plot_combined_summary(ft1_runs, ft3_runs, args.out_dir) + + print(f"\n[plot] All figures written to: {args.out_dir}") + + +if __name__ == "__main__": + main() diff --git a/examples/opf/finetune/plot_sample_efficiency.py b/examples/opf/finetune/plot_sample_efficiency.py new file mode 100755 index 000000000..aa9b0c11f --- /dev/null +++ b/examples/opf/finetune/plot_sample_efficiency.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +"""Aggregate OPF FT1 + FT3 single-method, single-N runs into sample-efficiency plots. + +Scans ``examples/opf/finetune/logs/`` for directories produced by +``train_opf_ft1_classify.py`` and ``train_opf_finetune.py``, pulls the test +metric from each run's ``results.json``, and produces one PNG per +(task, architecture, metric) showing the metric versus number of training +samples for each fine-tuning method (full, partial, head_only, scratch). + +Outputs land in ``examples/opf/finetune/plots/``. +""" +from __future__ import annotations + +import argparse +import json +import re +from pathlib import Path +from collections import defaultdict + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt + +# log directory naming conventions used by the training scripts +FT1_RE = re.compile( + r"^FT1_feasibility_(?P[A-Za-z0-9]+)_(?Pfull|partial|head_only)" + r"(?P_scratch)?_n(?P\d+)$" +) +FT3_RE = re.compile( + r"^finetune_FT3_contingency_(?P[A-Za-z0-9]+)_(?Pfull|partial|head_only)" + r"(?P_scratch)?_n(?P\d+)$" +) + +# which metric to plot per task (key inside results.json["test_metrics"]) +FT1_METRICS = ["bce", "accuracy", "f1", "auc_roc"] +FT3_METRICS = ["overall_mse"] + +METHOD_ORDER = ["full", "partial", "head_only", "scratch"] +METHOD_STYLE = { + "full": {"color": "#1f77b4", "marker": "o"}, + "partial": {"color": "#ff7f0e", "marker": "s"}, + "head_only": {"color": "#2ca02c", "marker": "^"}, + "scratch": {"color": "#d62728", "marker": "x"}, +} + + +def collect(logs_dir: Path) -> dict: + """Return nested dict: results[task][arch][method][n] = metrics dict.""" + results: dict = defaultdict(lambda: defaultdict(lambda: defaultdict(dict))) + for d in sorted(logs_dir.iterdir()): + if not d.is_dir(): + continue + m1 = FT1_RE.match(d.name) + m3 = FT3_RE.match(d.name) + if m1: + task = "FT1" + arch = m1["arch"] + method = "scratch" if m1["scratch"] else m1["regime"] + # LOG_NAME is now keyed on requested N directly. + n = int(m1["total"]) + elif m3: + task = "FT3" + arch = m3["arch"] + method = "scratch" if m3["scratch"] else m3["regime"] + n = int(m3["n"]) + else: + continue + + results_file = d / "results.json" + if not results_file.is_file(): + continue + try: + payload = json.loads(results_file.read_text()) + except (OSError, json.JSONDecodeError): + continue + metrics = payload.get("test_metrics", {}) + if not metrics: + continue + results[task][arch][method][n] = metrics + return results + + +def plot_one(task: str, arch: str, metric: str, + per_method: dict, out_path: Path) -> None: + """Draw a single sample-efficiency curve and write to ``out_path``.""" + fig, ax = plt.subplots(figsize=(7, 5)) + plotted_any = False + for method in METHOD_ORDER: + run_map = per_method.get(method, {}) + if not run_map: + continue + ns = sorted(run_map) + ys = [run_map[n].get(metric) for n in ns] + pairs = [(n, y) for n, y in zip(ns, ys) if y is not None] + if not pairs: + continue + xs, ys = zip(*pairs) + style = METHOD_STYLE[method] + ax.plot(xs, ys, label=method, linewidth=1.6, + marker=style["marker"], color=style["color"]) + plotted_any = True + + if not plotted_any: + plt.close(fig) + return + + ax.set_xscale("log") + if metric in {"bce", "overall_mse"}: + ax.set_yscale("log") + ax.set_xlabel("Number of fine-tuning samples (N)") + ax.set_ylabel(f"Test {metric}") + ax.set_title(f"{task} {arch} — sample efficiency ({metric})") + ax.grid(True, which="both", linestyle=":", linewidth=0.5) + ax.legend(title="Fine-tuning regime") + fig.tight_layout() + fig.savefig(out_path, dpi=150) + plt.close(fig) + print(f"[plot] {out_path}") + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--logs", type=Path, + default=Path(__file__).parent / "logs", + help="Path to logs directory") + ap.add_argument("--out", type=Path, + default=Path(__file__).parent / "plots", + help="Output directory for PNGs and aggregated CSV") + args = ap.parse_args() + + args.out.mkdir(parents=True, exist_ok=True) + results = collect(args.logs) + + # write a flat CSV summary for downstream use in the manuscript + summary_csv = args.out / "summary.csv" + rows = ["task,arch,method,n,metric,value"] + for task, archs in results.items(): + metric_keys = FT1_METRICS if task == "FT1" else FT3_METRICS + for arch, methods in archs.items(): + for metric in metric_keys: + per_method = {m: methods.get(m, {}) for m in METHOD_ORDER} + out_path = args.out / f"{task}_{arch}_{metric}.png" + plot_one(task, arch, metric, per_method, out_path) + for method, ns in methods.items(): + for n, mvals in sorted(ns.items()): + for k, v in mvals.items(): + if isinstance(v, (int, float)): + rows.append(f"{task},{arch},{method},{n},{k},{v}") + summary_csv.write_text("\n".join(rows) + "\n") + print(f"[csv] {summary_csv} ({len(rows) - 1} rows)") + + +if __name__ == "__main__": + main() diff --git a/examples/opf/finetune/preprocess_ft1_ft3_loginnode.sh b/examples/opf/finetune/preprocess_ft1_ft3_loginnode.sh new file mode 100644 index 000000000..6b904769c --- /dev/null +++ b/examples/opf/finetune/preprocess_ft1_ft3_loginnode.sh @@ -0,0 +1,73 @@ +#!/bin/bash +# ============================================================================= +# preprocess_ft1_ft3_loginnode.sh — CPU-only data preprocessing for FT1+FT3 +# +# Runs all preprocessing steps on the login node (no GPU, no srun): +# Step 1: Serialize FT3 contingency data (shared by HeteroSAGE and HeteroHEAT) +# Step 2: Generate FT1 feasibility dataset from FT3 data +# +# Usage (from examples/opf/finetune/): +# bash preprocess_ft1_ft3_loginnode.sh +# ============================================================================= +set -euo pipefail + +HYDRAGNN_ROOT=${HYDRAGNN_ROOT:-/lustre/orion/lrn078/proj-shared/HydraGNN} +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +DATASET_DIR="$HYDRAGNN_ROOT/examples/opf/dataset" + +# ───────────────────────────────────────────────────────────────────────────── +# Environment +# ───────────────────────────────────────────────────────────────────────────── +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:${PYTHONPATH:-} + +cd "$FT_DIR" + +# ───────────────────────────────────────────────────────────────────────────── +# Step 1: FT3 contingency data — HeteroSAGE +# ───────────────────────────────────────────────────────────────────────────── +echo "========================================================" +echo " Step 1: FT3 contingency preprocessing (HeteroSAGE)" +echo "========================================================" +python -u "$HYDRAGNN_ROOT/examples/opf/train_opf_solution_heterogeneous.py" \ + --inputfile "$FT_DIR/FT3_contingency/config_HeteroSAGE_full.json" \ + --hdf5 \ + --preonly \ + --case_name pglib_opf_case118_ieee \ + --num_groups 20 \ + --modelname FT3_contingency_data \ + --data_root "$DATASET_DIR" \ + --topological_perturbations +echo "Step 1 done." + +# ───────────────────────────────────────────────────────────────────────────── +# Step 2: FT1 feasibility dataset (from FT3 SAGE data) +# ───────────────────────────────────────────────────────────────────────────── +echo "" +echo "========================================================" +echo " Step 2: FT1 feasibility data generation" +echo "========================================================" +python -u "$FT_DIR/generate_infeasible_samples.py" \ + --src_dir "$DATASET_DIR/FT3_contingency_data" \ + --out_dir "$DATASET_DIR/FT1_feasibility_data" \ + --overload_factor 6.0 \ + --max_samples 5000 \ + --seed 42 +echo "Step 2 done." + +echo "" +echo "========================================================" +echo " All preprocessing complete." +echo " FT3 SAGE : $DATASET_DIR/FT3_contingency_data" +echo " FT3 HEAT : (reuses FT3_contingency_data)" +echo " FT1 data : $DATASET_DIR/FT1_feasibility_data" +echo "========================================================" diff --git a/examples/opf/finetune/submit_ft1_data_efficiency_sweep.sh b/examples/opf/finetune/submit_ft1_data_efficiency_sweep.sh new file mode 100644 index 000000000..9e1146891 --- /dev/null +++ b/examples/opf/finetune/submit_ft1_data_efficiency_sweep.sh @@ -0,0 +1,90 @@ +#!/bin/bash +# ============================================================================= +# submit_ft1_data_efficiency_sweep.sh +# +# Data-efficiency experiment: compare fine-tuning a pre-trained graph +# foundation model vs. training from scratch (random init) on FT1 feasibility +# classification, for progressively increasing numbers of training samples. +# +# Experiment grid: +# sample_sizes : SAMPLE_SIZES array below (training samples per class, +# i.e. total = 2 × N balanced feasible/infeasible) +# methods : full FT, partial FT, head_only FT, from-scratch baseline +# architectures : HeteroSAGE, HeteroHEAT +# +# Requirements: +# - FT1 feasibility dataset generated with enough samples +# (run generate_infeasible_samples.py --max_samples 100000 first) +# - Pretrained models in examples/opf/pretrained_models/ +# +# Usage (from examples/opf/finetune/): +# bash submit_ft1_data_efficiency_sweep.sh [--dry-run] +# ============================================================================= +set -euo pipefail + +# ───────────────────────────────────────────────────────────────────────────── +# Configuration +# ───────────────────────────────────────────────────────────────────────────── +HYDRAGNN_ROOT=${HYDRAGNN_ROOT:-/lustre/orion/lrn078/proj-shared/HydraGNN} +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +SLURM_OUT_DIR="$HYDRAGNN_ROOT" + +PROJECT=${PROJECT:-eng164} +PARTITION=${PARTITION:-batch} +# Nodes per individual training run; total allocation = 4 × N_PER_RUN +N_PER_RUN=${N_PER_RUN:-8} +N_NODES=$((N_PER_RUN * 4)) +WALL_TIME=${WALL_TIME:-02:00:00} + +# Progressively increasing number of feasible training samples. +# Total training samples = 2 × N (balanced feasible + infeasible). +SAMPLE_SIZES=(100 500 1000 2500 5000 10000 25000 50000 100000) + +DRY_RUN=false +for arg in "$@"; do + [[ "$arg" == "--dry-run" ]] && DRY_RUN=true +done + +# ───────────────────────────────────────────────────────────────────────────── +# Submission helper +# ───────────────────────────────────────────────────────────────────────────── +_sbatch() { + local label="$1"; shift + if $DRY_RUN; then + echo "[dry-run] sbatch $*" + else + local id + id=$(sbatch "$@" | awk '{print $NF}') + echo "[submit] $label → job $id" + fi +} + +echo "========================================================" +echo " FT1 data-efficiency sweep (multi-method jobs)" +echo " Account : $PROJECT" +echo " Partition : $PARTITION" +echo " Nodes/run : $N_PER_RUN (total per job: $N_NODES)" +echo " Wall time : $WALL_TIME" +echo " Sample sizes : ${SAMPLE_SIZES[*]}" +echo " Methods : full / partial / head_only / scratch (packed per job)" +echo "========================================================" + +for ARCH in HeteroSAGE HeteroHEAT; do + for N in "${SAMPLE_SIZES[@]}"; do + TOTAL=$((N * 2)) + LABEL="FT1_${ARCH}_multi_n${N}" + echo "--- $LABEL (total train=${TOTAL}, 4 methods × ${N_PER_RUN} nodes) ---" + _sbatch "$LABEL" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name="$LABEL" \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --error="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=$ARCH,N_PER_RUN=$N_PER_RUN,MAX_TRAIN_SAMPLES=$TOTAL \ + "$FT_DIR/FT1_feasibility_classification/job-frontier-FT1-multi.sh" + done +done + +echo "" +echo "All jobs submitted." +echo "Total jobs: $((${#SAMPLE_SIZES[@]} * 2)) (${#SAMPLE_SIZES[@]} sizes × 2 archs, 4 methods packed per job)" diff --git a/examples/opf/finetune/submit_ft1_ft3_debug.sh b/examples/opf/finetune/submit_ft1_ft3_debug.sh new file mode 100644 index 000000000..7519d5d0c --- /dev/null +++ b/examples/opf/finetune/submit_ft1_ft3_debug.sh @@ -0,0 +1,112 @@ +#!/bin/bash +# ============================================================================= +# submit_ft1_ft3_debug.sh — Debug-queue submission for FT1 + FT3 +# +# Submits fine-tuning jobs for the two best HPO architectures (HeteroSAGE +# and HeteroHEAT) under FT1 (feasibility classification) and FT3 +# (contingency) strategies using the Frontier debug queue. +# +# Assumptions: +# - Preprocessed datasets already exist (PHASES=train only). +# - Pretrained models are in examples/opf/pretrained_models/. +# +# Usage (from examples/opf/finetune/): +# bash submit_ft1_ft3_debug.sh [--dry-run] +# +# Flags: +# --dry-run Print sbatch commands without submitting. +# ============================================================================= +set -euo pipefail + +# ───────────────────────────────────────────────────────────────────────────── +# Configuration +# ───────────────────────────────────────────────────────────────────────────── +HYDRAGNN_ROOT=${HYDRAGNN_ROOT:-/lustre/orion/lrn078/proj-shared/HydraGNN} +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +SLURM_OUT_DIR="$HYDRAGNN_ROOT" + +PROJECT=eng164 +PARTITION=debug +N_NODES=2 # debug queue max +WALL_TIME=00:30:00 # debug queue max + +DRY_RUN=false +for arg in "$@"; do + [[ "$arg" == "--dry-run" ]] && DRY_RUN=true +done + +# ───────────────────────────────────────────────────────────────────────────── +# Submission helper +# ───────────────────────────────────────────────────────────────────────────── +_sbatch() { + local label="$1" + shift + if $DRY_RUN; then + echo "[dry-run] sbatch $*" + else + local id + id=$(sbatch "$@" | awk '{print $NF}') + echo "[submit] $label → job $id" + fi +} + +echo "========================================================" +echo " FT1 + FT3 debug submissions" +echo " Account : $PROJECT" +echo " Partition : $PARTITION" +echo " Nodes : $N_NODES" +echo " Wall time : $WALL_TIME" +echo "========================================================" + +# ───────────────────────────────────────────────────────────────────────────── +# FT1 — Feasibility classification (HeteroSAGE and HeteroHEAT) +# Assumes FT1 feasibility dataset already generated (PHASES=train). +# ───────────────────────────────────────────────────────────────────────────── +echo "" +echo "--- FT1: HeteroSAGE ---" +_sbatch "FT1_HeteroSAGE" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name=FT1-SAGE-dbg \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/FT1-SAGE-dbg-%j.out" \ + --error="$SLURM_OUT_DIR/FT1-SAGE-dbg-%j.out" \ + --export=ALL,FT_REGIME=full,PHASES=train \ + "$FT_DIR/FT1_feasibility_classification/job-frontier-FT1-HeteroSAGE.sh" + +echo "--- FT1: HeteroHEAT ---" +_sbatch "FT1_HeteroHEAT" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name=FT1-HEAT-dbg \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/FT1-HEAT-dbg-%j.out" \ + --error="$SLURM_OUT_DIR/FT1-HEAT-dbg-%j.out" \ + --export=ALL,FT_REGIME=full,PHASES=train \ + "$FT_DIR/FT1_feasibility_classification/job-frontier-FT1-HeteroHEAT.sh" + +# ───────────────────────────────────────────────────────────────────────────── +# FT3 — Contingency fine-tuning (HeteroSAGE and HeteroHEAT) +# Assumes FT3 contingency dataset already preprocessed (PHASES=train). +# ───────────────────────────────────────────────────────────────────────────── +echo "" +echo "--- FT3: HeteroSAGE ---" +_sbatch "FT3_HeteroSAGE" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name=FT3-SAGE-dbg \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/FT3-SAGE-dbg-%j.out" \ + --error="$SLURM_OUT_DIR/FT3-SAGE-dbg-%j.out" \ + --export=ALL,FT_STRATEGY=FT3_contingency,FT_ARCH=HeteroSAGE,FT_REGIME=full,PHASES=train \ + "$FT_DIR/FT3_contingency/job-frontier-HeteroSAGE.sh" + +echo "--- FT3: HeteroHEAT ---" +_sbatch "FT3_HeteroHEAT" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name=FT3-HEAT-dbg \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/FT3-HEAT-dbg-%j.out" \ + --error="$SLURM_OUT_DIR/FT3-HEAT-dbg-%j.out" \ + --export=ALL,FT_STRATEGY=FT3_contingency,FT_ARCH=HeteroHEAT,FT_REGIME=full,PHASES=train \ + "$FT_DIR/FT3_contingency/job-frontier-HeteroHEAT.sh" + +echo "" +echo "All jobs submitted." diff --git a/examples/opf/finetune/submit_ft1_ft3_jobs.sh b/examples/opf/finetune/submit_ft1_ft3_jobs.sh new file mode 100644 index 000000000..0e6e90cca --- /dev/null +++ b/examples/opf/finetune/submit_ft1_ft3_jobs.sh @@ -0,0 +1,319 @@ +#!/bin/bash +# ============================================================================= +# submit_ft1_ft3_jobs.sh — Frontier Slurm orchestration for FT1 + FT3 +# +# Submits all experiments as Slurm jobs with proper dependencies so that: +# +# Stage 0: Preprocess FT3 data (download + serialise, two archs in parallel) +# Stage 0: Generate FT1 feasibility data (after FT3 SAGE data is ready) +# Stage 1: Run 16 training jobs (8 FT1 + 8 FT3) in parallel after data is ready +# Stage 2: Collect results & generate plots after all 16 jobs complete +# +# Total submitted: 3 (data) + 16 (training) + 1 (collect+plot) = 20 jobs +# +# Usage (from examples/opf/finetune/): +# bash submit_ft1_ft3_jobs.sh [--dry-run] +# +# Flags: +# --dry-run Print all sbatch commands but do not submit them. +# +# Requirements: +# - HydraGNN pretrained models in examples/opf/pretrained_models/ +# - HYDRAGNN_ROOT, PROJECT, VENV variables set below or in environment +# ============================================================================= +set -euo pipefail + +# ───────────────────────────────────────────────────────────────────────────── +# Configuration — edit for your allocation +# ───────────────────────────────────────────────────────────────────────────── +HYDRAGNN_ROOT=${HYDRAGNN_ROOT:-/lustre/orion/lrn078/proj-shared/HydraGNN} +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +LOG_DIR="$HYDRAGNN_ROOT/logs" +PROJECT=${PROJECT:-eng164} +PARTITION=${PARTITION:-batch} +N_TRAIN_NODES=${N_TRAIN_NODES:-8} # nodes per training job +N_DATA_NODES=${N_DATA_NODES:-1} # nodes for preprocessing +TRAIN_TIME=${TRAIN_TIME:-04:00:00} # wall time per training job +DATA_TIME=${DATA_TIME:-02:00:00} # wall time for data jobs +SLURM_OUT_DIR=${SLURM_OUT_DIR:-/lustre/orion/lrn078/proj-shared/HydraGNN} + +DRY_RUN=false +for arg in "$@"; do + [[ "$arg" == "--dry-run" ]] && DRY_RUN=true +done + +# ───────────────────────────────────────────────────────────────────────────── +# Submission helper +# ───────────────────────────────────────────────────────────────────────────── +declare -a ALL_TRAINING_JOB_IDS=() + +_sbatch() { + local label="$1" + shift + if $DRY_RUN; then + echo "[dry-run] sbatch $*" + echo "DRY_$label" + else + local id + id=$(sbatch "$@" | awk '{print $NF}') + echo "[submit] $label → job $id" + echo "$id" + fi +} + +# ───────────────────────────────────────────────────────────────────────────── +# Stage 0a — Preprocess FT3 data (HeteroSAGE and HeteroHEAT in parallel) +# Phase 1 of train_opf_solution_heterogeneous.py --preonly +# ───────────────────────────────────────────────────────────────────────────── +echo "========================================================" +echo " Stage 0a: FT3 data preprocessing" +echo "========================================================" + +JOB_FT3_DATA_SAGE=$(_sbatch "FT3_data_HeteroSAGE" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name=FT3-data-SAGE \ + --nodes=$N_DATA_NODES --time=$DATA_TIME \ + --output="$SLURM_OUT_DIR/FT3-data-SAGE-%j.out" \ + --error="$SLURM_OUT_DIR/FT3-data-SAGE-%j.out" \ + --export=ALL,FT_STRATEGY=FT3_contingency,FT_ARCH=HeteroSAGE,FT_REGIME=full,PHASES=preonly \ + "$FT_DIR/FT3_contingency/job-frontier-HeteroSAGE.sh") + +# FT3 HeteroHEAT reuses the same dataset as HeteroSAGE — no separate data job needed. +JOB_FT3_DATA_HEAT=$JOB_FT3_DATA_SAGE + +# ───────────────────────────────────────────────────────────────────────────── +# Stage 0b — Generate FT1 feasibility dataset +# Depends on FT3_SAGE data (used as the source of feasible samples) +# ───────────────────────────────────────────────────────────────────────────── +echo "" +echo "========================================================" +echo " Stage 0b: FT1 feasibility data generation" +echo "========================================================" + +# Inline job script for generate_infeasible_samples.py +read -r -d '' FT1_DATA_SCRIPT <<'SCRIPT' || true +#!/bin/bash +#SBATCH -A __PROJECT__ +#SBATCH -J FT1-data +#SBATCH -o __SLURM_OUT__/FT1-data-%j.out +#SBATCH -e __SLURM_OUT__/FT1-data-%j.out +#SBATCH -t __DATA_TIME__ +#SBATCH -p __PARTITION__ +#SBATCH -N 1 + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate __VENV__ + +HYDRAGNN_ROOT=__HYDRAGNN_ROOT__ +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +DATASET_DIR="$HYDRAGNN_ROOT/examples/opf/dataset" + +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +SRC_DIR="$DATASET_DIR/FT3_contingency_data" +OUT_DIR="$DATASET_DIR/FT1_feasibility_data" + +echo "[FT1 data] Generating infeasible samples..." +echo " src : $SRC_DIR" +echo " out : $OUT_DIR" + +python -u "$FT_DIR/generate_infeasible_samples.py" \ + --src_dir "$SRC_DIR" \ + --out_dir "$OUT_DIR" \ + --overload_factor 6.0 \ + --max_samples 5000 \ + --seed 42 + +echo "[FT1 data] Done." +SCRIPT + +# Materialise the inline script to a temp file +FT1_DATA_SCRIPT_PATH="$FT_DIR/_tmp_ft1_data_job.sh" +echo "$FT1_DATA_SCRIPT" \ + | sed "s|__PROJECT__|$PROJECT|g" \ + | sed "s|__SLURM_OUT__|$SLURM_OUT_DIR|g" \ + | sed "s|__DATA_TIME__|$DATA_TIME|g" \ + | sed "s|__PARTITION__|$PARTITION|g" \ + | sed "s|__HYDRAGNN_ROOT__|$HYDRAGNN_ROOT|g" \ + | sed "s|__VENV__|/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv|g" \ + > "$FT1_DATA_SCRIPT_PATH" +chmod +x "$FT1_DATA_SCRIPT_PATH" + +JOB_FT1_DATA=$(_sbatch "FT1_data" \ + --dependency="afterok:$JOB_FT3_DATA_SAGE" \ + "$FT1_DATA_SCRIPT_PATH") + +# ───────────────────────────────────────────────────────────────────────────── +# Stage 1 — Training jobs +# ───────────────────────────────────────────────────────────────────────────── +echo "" +echo "========================================================" +echo " Stage 1: Training jobs (FT1 + FT3, all regimes)" +echo "========================================================" + +# Helper: submit a single FT1 training job +_submit_ft1() { + local arch="$1" + local regime="$2" + local no_pretrained="${3:-false}" + local label="FT1_${arch}_${regime}" + local dep="afterok:$JOB_FT1_DATA" + local extra_flags="" + + if [[ "$no_pretrained" == "true" ]]; then + label="FT1_${arch}_baseline" + extra_flags="--export=ALL,FT_REGIME=full,NO_PRETRAINED=1" + else + extra_flags="--export=ALL,FT_REGIME=$regime" + fi + + local job_script="$FT_DIR/FT1_feasibility_classification/job-frontier-FT1-${arch}.sh" + + local id + id=$(_sbatch "$label" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name="$label" \ + --nodes=$N_TRAIN_NODES --time=$TRAIN_TIME \ + --output="$SLURM_OUT_DIR/${label}-%j.out" \ + --error="$SLURM_OUT_DIR/${label}-%j.out" \ + --dependency="$dep" \ + $extra_flags \ + "$job_script") + ALL_TRAINING_JOB_IDS+=("$id") +} + +# Helper: submit a single FT3 training job +_submit_ft3() { + local arch="$1" + local regime="$2" + local no_pretrained="${3:-false}" + local label="FT3_${arch}_${regime}" + local dep="afterok:$JOB_FT3_DATA_SAGE" + local export_str="FT_STRATEGY=FT3_contingency,FT_ARCH=$arch,FT_REGIME=$regime,PHASES=train" + + if [[ "$no_pretrained" == "true" ]]; then + label="FT3_${arch}_baseline" + export_str="${export_str},NO_PRETRAINED=1" + fi + + local job_script + if [[ "$arch" == "HeteroSAGE" ]]; then + job_script="$FT_DIR/FT3_contingency/job-frontier-HeteroSAGE.sh" + else + job_script="$FT_DIR/FT3_contingency/job-frontier-HeteroHEAT.sh" + fi + + local id + id=$(_sbatch "$label" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name="$label" \ + --nodes=$N_TRAIN_NODES --time=$TRAIN_TIME \ + --output="$SLURM_OUT_DIR/${label}-%j.out" \ + --error="$SLURM_OUT_DIR/${label}-%j.out" \ + --dependency="$dep" \ + --export="ALL,$export_str" \ + "$job_script") + ALL_TRAINING_JOB_IDS+=("$id") +} + +# ── FT1 × {HeteroSAGE, HeteroHEAT} × {head_only, partial, full, baseline} ── +for ARCH in HeteroSAGE HeteroHEAT; do + for REGIME in head_only partial full; do + _submit_ft1 "$ARCH" "$REGIME" false + done + _submit_ft1 "$ARCH" "full" true # baseline (random init) +done + +# ── FT3 × {HeteroSAGE, HeteroHEAT} × {head_only, partial, full, baseline} ── +for ARCH in HeteroSAGE HeteroHEAT; do + for REGIME in head_only partial full; do + _submit_ft3 "$ARCH" "$REGIME" false + done + _submit_ft3 "$ARCH" "full" true # baseline (random init) +done + +echo "" +echo "Submitted ${#ALL_TRAINING_JOB_IDS[@]} training jobs:" +printf ' %s\n' "${ALL_TRAINING_JOB_IDS[@]}" + +# ───────────────────────────────────────────────────────────────────────────── +# Stage 2 — Collect results and generate plots +# ───────────────────────────────────────────────────────────────────────────── +echo "" +echo "========================================================" +echo " Stage 2: collect_results.py + plot_ft_results.py" +echo "========================================================" + +# Build afterok dependency on ALL training jobs +DEP_ALL_TRAIN=$(IFS=:; echo "afterok:${ALL_TRAINING_JOB_IDS[*]}") + +read -r -d '' COLLECT_SCRIPT <<'SCRIPT' || true +#!/bin/bash +#SBATCH -A __PROJECT__ +#SBATCH -J FT-collect +#SBATCH -o __SLURM_OUT__/FT-collect-%j.out +#SBATCH -e __SLURM_OUT__/FT-collect-%j.out +#SBATCH -t 00:30:00 +#SBATCH -p __PARTITION__ +#SBATCH -N 1 + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate __VENV__ + +HYDRAGNN_ROOT=__HYDRAGNN_ROOT__ +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +LOGS_ROOT="$HYDRAGNN_ROOT/logs" +OUT_DIR="$FT_DIR/results" + +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +echo "[Stage 2] Collecting results..." +python -u "$FT_DIR/collect_results.py" \ + --logs_root "$LOGS_ROOT" \ + --out_dir "$OUT_DIR" + +echo "[Stage 2] Generating plots..." +python -u "$FT_DIR/plot_ft_results.py" \ + --summary "$OUT_DIR/ft1_ft3_summary.json" \ + --out_dir "$OUT_DIR/figures" + +echo "[Stage 2] Done. Figures in: $OUT_DIR/figures/" +SCRIPT + +COLLECT_SCRIPT_PATH="$FT_DIR/_tmp_collect_job.sh" +echo "$COLLECT_SCRIPT" \ + | sed "s|__PROJECT__|$PROJECT|g" \ + | sed "s|__SLURM_OUT__|$SLURM_OUT_DIR|g" \ + | sed "s|__PARTITION__|$PARTITION|g" \ + | sed "s|__HYDRAGNN_ROOT__|$HYDRAGNN_ROOT|g" \ + | sed "s|__VENV__|/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv|g" \ + > "$COLLECT_SCRIPT_PATH" +chmod +x "$COLLECT_SCRIPT_PATH" + +JOB_COLLECT=$(_sbatch "collect_plot" \ + --dependency="$DEP_ALL_TRAIN" \ + "$COLLECT_SCRIPT_PATH") + +# ───────────────────────────────────────────────────────────────────────────── +# Summary +# ───────────────────────────────────────────────────────────────────────────── +echo "" +echo "========================================================" +echo " Submission complete" +echo "========================================================" +echo " FT3 data prep : $JOB_FT3_DATA_SAGE (SAGE), $JOB_FT3_DATA_HEAT (HEAT)" +echo " FT1 data gen : $JOB_FT1_DATA" +echo " Training jobs : ${ALL_TRAINING_JOB_IDS[*]}" +echo " Collect+plot : $JOB_COLLECT" +echo "" +echo "Monitor with:" +echo " squeue -u \$USER" +echo " tail -f $SLURM_OUT_DIR/FT-collect-.out" +echo "" +echo "After completion, results are in:" +echo " $HYDRAGNN_ROOT/examples/opf/finetune/results/" + +# Cleanup temp scripts (only if we actually submitted) +if ! $DRY_RUN; then + rm -f "$FT1_DATA_SCRIPT_PATH" "$COLLECT_SCRIPT_PATH" +fi diff --git a/examples/opf/finetune/submit_ft3_data_efficiency_sweep.sh b/examples/opf/finetune/submit_ft3_data_efficiency_sweep.sh new file mode 100644 index 000000000..d77c5c31e --- /dev/null +++ b/examples/opf/finetune/submit_ft3_data_efficiency_sweep.sh @@ -0,0 +1,87 @@ +#!/bin/bash +# ============================================================================= +# submit_ft3_data_efficiency_sweep.sh +# +# Data-efficiency experiment for FT3 (N-1 contingency OPF regression): +# compare fine-tuning a pre-trained graph foundation model vs. training from +# scratch for progressively increasing numbers of training samples. +# +# Experiment grid: +# sample_sizes : SAMPLE_SIZES array below +# methods : full FT, partial FT, head_only FT, from-scratch baseline +# architectures : HeteroSAGE, HeteroHEAT +# +# Requirements: +# - FT3 contingency dataset already serialised at +# examples/opf/dataset/FT3_contingency_data.h5 +# - Pretrained models in examples/opf/pretrained_models/ +# +# Usage (from examples/opf/finetune/): +# bash submit_ft3_data_efficiency_sweep.sh [--dry-run] +# ============================================================================= +set -euo pipefail + +# ───────────────────────────────────────────────────────────────────────────── +# Configuration +# ───────────────────────────────────────────────────────────────────────────── +HYDRAGNN_ROOT=${HYDRAGNN_ROOT:-/lustre/orion/lrn078/proj-shared/HydraGNN} +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +SLURM_OUT_DIR="$HYDRAGNN_ROOT" + +PROJECT=${PROJECT:-eng164} +PARTITION=${PARTITION:-batch} +# Nodes per individual training run; total allocation = 4 × N_PER_RUN +N_PER_RUN=${N_PER_RUN:-8} +N_NODES=$((N_PER_RUN * 4)) +WALL_TIME=${WALL_TIME:-02:00:00} + +# FT3 has 270k train samples available (full FT3_contingency_data.h5 trainset) +SAMPLE_SIZES=(100 500 1000 2500 5000 10000 25000 50000 100000 270000) + +DRY_RUN=false +for arg in "$@"; do + [[ "$arg" == "--dry-run" ]] && DRY_RUN=true +done + +# ───────────────────────────────────────────────────────────────────────────── +# Submission helper +# ───────────────────────────────────────────────────────────────────────────── +_sbatch() { + local label="$1"; shift + if $DRY_RUN; then + echo "[dry-run] sbatch $*" + else + local id + id=$(sbatch "$@" | awk '{print $NF}') + echo "[submit] $label → job $id" + fi +} + +echo "========================================================" +echo " FT3 data-efficiency sweep (multi-method jobs)" +echo " Account : $PROJECT" +echo " Partition : $PARTITION" +echo " Nodes/run : $N_PER_RUN (total per job: $N_NODES)" +echo " Wall time : $WALL_TIME" +echo " Sample sizes : ${SAMPLE_SIZES[*]}" +echo " Methods : full / partial / head_only / scratch (packed per job)" +echo "========================================================" + +for ARCH in HeteroSAGE HeteroHEAT; do + for N in "${SAMPLE_SIZES[@]}"; do + LABEL="FT3_${ARCH}_multi_n${N}" + echo "--- $LABEL (4 methods × ${N_PER_RUN} nodes) ---" + _sbatch "$LABEL" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name="$LABEL" \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --error="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=$ARCH,N_PER_RUN=$N_PER_RUN,MAX_TRAIN_SAMPLES=$N \ + "$FT_DIR/job-frontier-FT3-multi.sh" + done +done + +echo "" +echo "All jobs submitted." +echo "Total jobs: $((${#SAMPLE_SIZES[@]} * 2)) (${#SAMPLE_SIZES[@]} sizes × 2 archs, 4 methods packed per job)" diff --git a/examples/opf/finetune/submit_ft_bundled.sh b/examples/opf/finetune/submit_ft_bundled.sh new file mode 100755 index 000000000..f2ed896ef --- /dev/null +++ b/examples/opf/finetune/submit_ft_bundled.sh @@ -0,0 +1,112 @@ +#!/bin/bash +# ============================================================================= +# submit_ft_bundled.sh — Submit all FT1 + FT3 data-efficiency sweeps as +# 8 bundled Slurm jobs (2 per arch × task). +# +# Each job chains sample sizes sequentially; within each size the 4 training +# methods (full / partial / head_only / scratch) run concurrently via srun. +# Split into small/large bundles to stay within the 2h wall-time cap. +# +# Jobs submitted (per arch, ×2 for HeteroSAGE + HeteroHEAT): +# FT1-small n = 100 500 1000 2500 5000 +# FT1-large n = 10000 25000 50000 100000 +# FT3-small n = 100 500 1000 2500 5000 +# FT3-large n = 10000 25000 50000 100000 270000 +# +# Usage (from examples/opf/finetune/): +# bash submit_ft_bundled.sh [--dry-run] +# ============================================================================= +set -euo pipefail + +HYDRAGNN_ROOT=${HYDRAGNN_ROOT:-/lustre/orion/lrn078/proj-shared/HydraGNN} +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +SLURM_OUT_DIR="$HYDRAGNN_ROOT" + +PROJECT=${PROJECT:-eng164} +PARTITION=${PARTITION:-batch} +N_PER_RUN=${N_PER_RUN:-8} +N_NODES=$((N_PER_RUN * 4)) # 32 total nodes per job +WALL_TIME=${WALL_TIME:-02:00:00} + +FT1_SMALL_SIZES="100 500 1000 2500 5000" +FT1_LARGE_SIZES="10000 25000 50000 100000" +FT3_SMALL_SIZES="100 500 1000 2500 5000" +FT3_LARGE_SIZES="10000 25000 50000 100000 270000" + +DRY_RUN=false +for arg in "$@"; do + [[ "$arg" == "--dry-run" ]] && DRY_RUN=true +done + +_sbatch() { + local label="$1"; shift + if $DRY_RUN; then + echo "[dry-run] sbatch $*" + else + local id + id=$(sbatch "$@" | awk '{print $NF}') + echo "[submit] $label → job $id" + fi +} + +echo "========================================================" +echo " FT bundled data-efficiency sweep" +echo " Account : $PROJECT" +echo " Partition : $PARTITION" +echo " Nodes : $N_NODES (${N_PER_RUN} per method × 4 methods)" +echo " Wall time : $WALL_TIME" +echo " Jobs : 8 (FT1+FT3 × small+large × HeteroSAGE+HeteroHEAT)" +echo "========================================================" + +for ARCH in HeteroSAGE HeteroHEAT; do + # ── FT1 small ──────────────────────────────────────────────────────── + LABEL="FT1-small-${ARCH}" + echo "--- $LABEL (${FT1_SMALL_SIZES}) ---" + _sbatch "$LABEL" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name="$LABEL" \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --error="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=$ARCH,N_PER_RUN=$N_PER_RUN,SAMPLE_SIZES="$FT1_SMALL_SIZES" \ + "$FT_DIR/job-frontier-FT1-bundled.sh" + + # ── FT1 large ──────────────────────────────────────────────────────── + LABEL="FT1-large-${ARCH}" + echo "--- $LABEL (${FT1_LARGE_SIZES}) ---" + _sbatch "$LABEL" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name="$LABEL" \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --error="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=$ARCH,N_PER_RUN=$N_PER_RUN,SAMPLE_SIZES="$FT1_LARGE_SIZES" \ + "$FT_DIR/job-frontier-FT1-bundled.sh" + + # ── FT3 small ──────────────────────────────────────────────────────── + LABEL="FT3-small-${ARCH}" + echo "--- $LABEL (${FT3_SMALL_SIZES}) ---" + _sbatch "$LABEL" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name="$LABEL" \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --error="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=$ARCH,N_PER_RUN=$N_PER_RUN,SAMPLE_SIZES="$FT3_SMALL_SIZES" \ + "$FT_DIR/job-frontier-FT3-bundled.sh" + + # ── FT3 large ──────────────────────────────────────────────────────── + LABEL="FT3-large-${ARCH}" + echo "--- $LABEL (${FT3_LARGE_SIZES}) ---" + _sbatch "$LABEL" \ + --account=$PROJECT --partition=$PARTITION \ + --job-name="$LABEL" \ + --nodes=$N_NODES --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --error="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=$ARCH,N_PER_RUN=$N_PER_RUN,SAMPLE_SIZES="$FT3_LARGE_SIZES" \ + "$FT_DIR/job-frontier-FT3-bundled.sh" +done + +echo "" +echo "Done. 8 jobs submitted." diff --git a/examples/opf/finetune/submit_ft_debug_chain.sh b/examples/opf/finetune/submit_ft_debug_chain.sh new file mode 100755 index 000000000..796c12a9f --- /dev/null +++ b/examples/opf/finetune/submit_ft_debug_chain.sh @@ -0,0 +1,51 @@ +#!/bin/bash +# Submit all 16 single-method FT jobs to debug QOS, chained via --dependency=afterany +# so only one runs at a time. +set -euo pipefail + +FT_DIR="$(cd "$(dirname "$0")" && pwd)" +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +DRY_RUN=false +for arg in "$@"; do + [[ "$arg" == "--dry-run" ]] && DRY_RUN=true +done + +PREV_JOB="" + +submit_one() { + local TASK="$1" ARCH="$2" METHOD="$3" + local LABEL="${TASK}-${ARCH}-${METHOD}" + local SCRIPT="${FT_DIR}/job-frontier-${TASK}-single-method.sh" + local DEP_OPT="" + [[ -n "$PREV_JOB" ]] && DEP_OPT="--dependency=afterany:${PREV_JOB}" + + if $DRY_RUN; then + echo "[dry-run] $LABEL dep=${PREV_JOB:-none}" + PREV_JOB="DRYRUN_${LABEL}" + return + fi + + local ID + ID=$(sbatch \ + --account=eng164 --partition=batch --qos=debug \ + --job-name="$LABEL" \ + --nodes=8 --time=02:00:00 \ + --output="${HYDRAGNN_ROOT}/${LABEL}-%j.out" \ + --error="${HYDRAGNN_ROOT}/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=${ARCH},FT_METHOD=${METHOD},N_PER_RUN=8 \ + ${DEP_OPT} \ + "$SCRIPT" | awk '{print $NF}') + echo "[submit] $LABEL → job $ID ${DEP_OPT}" + PREV_JOB="$ID" +} + +for ARCH in HeteroSAGE HeteroHEAT; do + for METHOD in full partial head_only scratch; do + submit_one FT1 "$ARCH" "$METHOD" + submit_one FT3 "$ARCH" "$METHOD" + done +done + +echo "" +echo "All 16 jobs submitted in debug QOS chain." diff --git a/examples/opf/finetune/submit_ft_debug_one.sh b/examples/opf/finetune/submit_ft_debug_one.sh new file mode 100755 index 000000000..9cf79d9d8 --- /dev/null +++ b/examples/opf/finetune/submit_ft_debug_one.sh @@ -0,0 +1,180 @@ +#!/bin/bash +# ============================================================================= +# submit_ft_debug_one.sh +# +# Submit exactly one FT bundled job to Frontier debug QOS. +# Debug is single-job-per-user on this system, so this script supports: +# - immediate submit (default) +# - wait-until-slot-opens then submit (--wait-slot) +# +# Usage: +# bash submit_ft_debug_one.sh --target FT1-small-HeteroSAGE +# bash submit_ft_debug_one.sh --target FT3-small-HeteroHEAT --wait-slot +# bash submit_ft_debug_one.sh --target FT1-large-HeteroSAGE --dry-run +# +# Valid targets: +# FT1-small-HeteroSAGE +# FT1-large-HeteroSAGE +# FT1-small-HeteroHEAT +# FT1-large-HeteroHEAT +# FT3-small-HeteroSAGE +# FT3-large-HeteroSAGE +# FT3-small-HeteroHEAT +# FT3-large-HeteroHEAT +# ============================================================================= +set -euo pipefail + +HYDRAGNN_ROOT=${HYDRAGNN_ROOT:-/lustre/orion/lrn078/proj-shared/HydraGNN} +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +SLURM_OUT_DIR="$HYDRAGNN_ROOT" + +PROJECT=${PROJECT:-eng164} +PARTITION=${PARTITION:-batch} +QOS=${QOS:-debug} +N_PER_RUN=${N_PER_RUN:-8} +N_NODES=$((N_PER_RUN * 4)) +WALL_TIME=${WALL_TIME:-02:00:00} + +TARGET="" +WAIT_SLOT=false +DRY_RUN=false + +while [[ $# -gt 0 ]]; do + case "$1" in + --target) + TARGET="${2:-}" + shift 2 + ;; + --wait-slot) + WAIT_SLOT=true + shift + ;; + --dry-run) + DRY_RUN=true + shift + ;; + -h|--help) + sed -n '1,42p' "$0" + exit 0 + ;; + *) + echo "Unknown argument: $1" + exit 1 + ;; + esac +done + +if [[ -z "$TARGET" ]]; then + echo "Error: --target is required." + exit 1 +fi + +SCRIPT_PATH="" +FT_ARCH="" +SAMPLE_SIZES="" + +case "$TARGET" in + FT1-small-HeteroSAGE) + SCRIPT_PATH="$FT_DIR/job-frontier-FT1-bundled.sh" + FT_ARCH="HeteroSAGE" + SAMPLE_SIZES="100 500 1000 2500 5000" + ;; + FT1-large-HeteroSAGE) + SCRIPT_PATH="$FT_DIR/job-frontier-FT1-bundled.sh" + FT_ARCH="HeteroSAGE" + SAMPLE_SIZES="10000 25000 50000 100000" + ;; + FT1-small-HeteroHEAT) + SCRIPT_PATH="$FT_DIR/job-frontier-FT1-bundled.sh" + FT_ARCH="HeteroHEAT" + SAMPLE_SIZES="100 500 1000 2500 5000" + ;; + FT1-large-HeteroHEAT) + SCRIPT_PATH="$FT_DIR/job-frontier-FT1-bundled.sh" + FT_ARCH="HeteroHEAT" + SAMPLE_SIZES="10000 25000 50000 100000" + ;; + FT3-small-HeteroSAGE) + SCRIPT_PATH="$FT_DIR/job-frontier-FT3-bundled.sh" + FT_ARCH="HeteroSAGE" + SAMPLE_SIZES="100 500 1000 2500 5000" + ;; + FT3-large-HeteroSAGE) + SCRIPT_PATH="$FT_DIR/job-frontier-FT3-bundled.sh" + FT_ARCH="HeteroSAGE" + SAMPLE_SIZES="10000 25000 50000 100000 270000" + ;; + FT3-small-HeteroHEAT) + SCRIPT_PATH="$FT_DIR/job-frontier-FT3-bundled.sh" + FT_ARCH="HeteroHEAT" + SAMPLE_SIZES="100 500 1000 2500 5000" + ;; + FT3-large-HeteroHEAT) + SCRIPT_PATH="$FT_DIR/job-frontier-FT3-bundled.sh" + FT_ARCH="HeteroHEAT" + SAMPLE_SIZES="10000 25000 50000 100000 270000" + ;; + *) + echo "Invalid --target: $TARGET" + exit 1 + ;; +esac + +_debug_jobs_active() { + squeue -u "$USER" -h -o "%q %T" | awk '$1=="debug" && $2 ~ /PENDING|RUNNING|CONFIGURING/{n++} END{print n+0}' +} + +if ! $DRY_RUN; then + if $WAIT_SLOT; then + echo "Waiting for a free debug QOS slot..." + while true; do + ACTIVE=$(_debug_jobs_active) + if [[ "$ACTIVE" -eq 0 ]]; then + echo "Debug slot is free." + break + fi + echo "Debug slot busy (active debug jobs: $ACTIVE). Retrying in 30s..." + sleep 30 + done + else + ACTIVE=$(_debug_jobs_active) + if [[ "$ACTIVE" -gt 0 ]]; then + echo "Debug slot currently busy (active debug jobs: $ACTIVE)." + echo "Use --wait-slot to auto-submit when free." + exit 2 + fi + fi +fi + +CMD=( + sbatch + --account="$PROJECT" + --partition="$PARTITION" + --qos="$QOS" + --job-name="$TARGET" + --nodes="$N_NODES" + --time="$WALL_TIME" + --output="$SLURM_OUT_DIR/${TARGET}-%j.out" + --error="$SLURM_OUT_DIR/${TARGET}-%j.out" + --export="ALL,FT_ARCH=$FT_ARCH,N_PER_RUN=$N_PER_RUN,SAMPLE_SIZES=$SAMPLE_SIZES" + "$SCRIPT_PATH" +) + +echo "========================================================" +echo " Submitting one debug FT bundle" +echo " Target : $TARGET" +echo " Account : $PROJECT" +echo " Partition : $PARTITION" +echo " QOS : $QOS" +echo " Nodes : $N_NODES" +echo " Wall time : $WALL_TIME" +echo " Script : $SCRIPT_PATH" +echo "========================================================" + +if $DRY_RUN; then + echo -n "[dry-run] " + printf '%q ' "${CMD[@]}" + echo +else + "${CMD[@]}" +fi diff --git a/examples/opf/finetune/submit_ft_single_method.sh b/examples/opf/finetune/submit_ft_single_method.sh new file mode 100755 index 000000000..fb9614ee8 --- /dev/null +++ b/examples/opf/finetune/submit_ft_single_method.sh @@ -0,0 +1,83 @@ +#!/bin/bash +# ============================================================================= +# submit_ft_single_method.sh +# +# Submits all FT1 + FT3 data-efficiency runs as 16 small jobs: +# 2 archs × 2 tasks × 4 methods = 16 jobs × 8 nodes each +# +# Each job chains all sample sizes sequentially for one method. +# Much easier to backfill than 8 × 32-node jobs. +# +# Usage: +# bash submit_ft_single_method.sh [--dry-run] +# ============================================================================= +set -euo pipefail + +HYDRAGNN_ROOT=${HYDRAGNN_ROOT:-/lustre/orion/lrn078/proj-shared/HydraGNN} +FT_DIR="$HYDRAGNN_ROOT/examples/opf/finetune" +SLURM_OUT_DIR="$HYDRAGNN_ROOT" + +PROJECT=${PROJECT:-eng164} +PARTITION=${PARTITION:-batch} +QOS=${QOS:-normal} +N_PER_RUN=${N_PER_RUN:-8} +WALL_TIME=${WALL_TIME:-02:00:00} + +DRY_RUN=false +for arg in "$@"; do + [[ "$arg" == "--dry-run" ]] && DRY_RUN=true +done + +_sbatch() { + local label="$1"; shift + if $DRY_RUN; then + echo "[dry-run] sbatch $*" + else + local id + id=$(sbatch "$@" | awk '{print $NF}') + echo "[submit] $label → job $id" + fi +} + +echo "========================================================" +echo " FT single-method data-efficiency sweep" +echo " Account : $PROJECT" +echo " Partition : $PARTITION" +echo " QOS : $QOS" +echo " Nodes : $N_PER_RUN per job (4x smaller than bundled)" +echo " Wall time : $WALL_TIME" +echo " Jobs : 16 (2 archs × 2 tasks × 4 methods)" +echo "========================================================" + +for ARCH in HeteroSAGE HeteroHEAT; do + for METHOD in full partial head_only scratch; do + + # ── FT1 ────────────────────────────────────────────────────────── + LABEL="FT1-${ARCH}-${METHOD}" + echo "--- $LABEL ---" + _sbatch "$LABEL" \ + --account=$PROJECT --partition=$PARTITION --qos=$QOS \ + --job-name="$LABEL" \ + --nodes=$N_PER_RUN --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --error="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=$ARCH,FT_METHOD=$METHOD,N_PER_RUN=$N_PER_RUN \ + "$FT_DIR/job-frontier-FT1-single-method.sh" + + # ── FT3 ────────────────────────────────────────────────────────── + LABEL="FT3-${ARCH}-${METHOD}" + echo "--- $LABEL ---" + _sbatch "$LABEL" \ + --account=$PROJECT --partition=$PARTITION --qos=$QOS \ + --job-name="$LABEL" \ + --nodes=$N_PER_RUN --time=$WALL_TIME \ + --output="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --error="$SLURM_OUT_DIR/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=$ARCH,FT_METHOD=$METHOD,N_PER_RUN=$N_PER_RUN \ + "$FT_DIR/job-frontier-FT3-single-method.sh" + + done +done + +echo "" +echo "Done. 16 jobs submitted (8 nodes each)." diff --git a/examples/opf/finetune/train_opf_finetune.py b/examples/opf/finetune/train_opf_finetune.py new file mode 100644 index 000000000..e62772904 --- /dev/null +++ b/examples/opf/finetune/train_opf_finetune.py @@ -0,0 +1,607 @@ +"""Fine-tuning script for HydraGNN OPF heterogeneous models. + +Loads a pretrained HeteroSAGE or HeteroHEAT checkpoint and fine-tunes it on +a target OPF dataset. Requires that the target dataset has already been +serialised to HDF5 using train_opf_solution_heterogeneous.py --preonly --hdf5. + +Three fine-tuning regimes are supported: + head_only -- freeze all input projectors and all conv/feature layers; + only the MLP prediction head is updated. + partial -- freeze input projectors and the first (num_conv_layers-1) + conv/feature layers; the last conv layer and the prediction + head are updated. + full -- all parameters are updated. + +Example (FT1, topology-specific, full fine-tuning with HeteroSAGE): + python train_opf_finetune.py \\ + --inputfile FT1_topology/config_HeteroSAGE_full.json \\ + --hdf5 \\ + --modelname FT1_case118_HeteroSAGE_full \\ + --pretrained_model_dir ../pretrained_models \\ + --pretrained_model_name HeteroSAGE_best \\ + --finetune_regime full \\ + --num_epoch 50 \\ + --learning_rate 5e-4 +""" + +import os +import sys +import json +import argparse +import shutil +import glob +import re + +# Make examples/opf importable +_OPF_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..") +sys.path.insert(0, _OPF_DIR) + +import torch +import torch.distributed as dist +from mpi4py import MPI + +import hydragnn +import hydragnn.utils.model as model_utils +from hydragnn.utils.model.model import load_existing_model +from hydragnn.utils.model import print_model +from hydragnn.utils.distributed import get_device +from hydragnn.utils.input_config_parsing.config_utils import update_config +from hydragnn.utils.input_config_parsing import save_config + +from hydragnn.utils.datasets.hdf5dataset import HDF5Dataset + +from ft_utils import EpochCSVWriter, evaluate_ft3, save_run_results + +from opf_solution_utils import ( + EdgeAttrDatasetAdapter, + OPFEnhancedModelWrapper, + NodeBatchAdapter, + NodeTargetDatasetAdapter, + OPFDomainLoss, + compute_pna_deg_for_hetero_dataset, + validate_voi_node_features, + info, + resolve_node_target_type as _resolve_node_target_type, +) + + +# --------------------------------------------------------------------------- +# Utility helpers +# --------------------------------------------------------------------------- + +def _to_jsonable(obj): + import numpy as np + if isinstance(obj, torch.Tensor): + return obj.item() if obj.numel() == 1 else obj.tolist() + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, np.generic): + return obj.item() + if isinstance(obj, dict): + return {k: _to_jsonable(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_to_jsonable(v) for v in obj] + return obj + + +def _resolve_edge_dim(config): + arch = config.get("NeuralNetwork", {}).get("Architecture", {}) + return arch.get("edge_dim", {"ac_line": 9, "transformer": 11}) + + +def apply_freeze_regime(model, regime: str): + """Freeze model parameters according to the requested fine-tuning regime. + + Must be called *after* loading pretrained weights but *before* creating the + optimizer so that frozen parameters are excluded from the parameter groups. + + Parameters + ---------- + model : nn.Module + The (not yet DDP-wrapped) model. + regime : str + One of 'full', 'partial', 'head_only'. + """ + if regime == "full": + return # nothing to freeze + + # Unwrap OPFEnhancedModelWrapper + inner = model + if isinstance(inner, OPFEnhancedModelWrapper): + inner = inner.model + + if regime == "head_only": + # Freeze everything except the prediction heads (heads_NN). + for name, param in inner.named_parameters(): + if "heads_NN" not in name: + param.requires_grad_(False) + n_trainable = sum(p.numel() for p in inner.parameters() if p.requires_grad) + info(f"[FT] head_only: {n_trainable:,} trainable parameters (prediction heads only)") + return + + if regime == "partial": + # Freeze input projectors and all but the last conv/feature layer pair. + for name, param in inner.named_parameters(): + if "input_projectors" in name or "node_embedders" in name: + param.requires_grad_(False) + + n_conv = len(inner.graph_convs) + for i, (conv, feat) in enumerate( + zip(inner.graph_convs, inner.feature_layers) + ): + if i < n_conv - 1: # freeze all but the last conv layer + for p in conv.parameters(): + p.requires_grad_(False) + for p in feat.parameters(): + p.requires_grad_(False) + + n_trainable = sum(p.numel() for p in inner.parameters() if p.requires_grad) + info( + f"[FT] partial: {n_trainable:,} trainable parameters " + f"(last conv layer + heads; {n_conv - 1} of {n_conv} conv layers frozen)" + ) + return + + raise ValueError( + f"Unknown finetune_regime '{regime}'. Choose from: full, partial, head_only." + ) + + +def load_pretrained_weights(model, pretrained_model_name: str, pretrained_model_dir: str): + """Load model weights from a pretrained checkpoint, discarding optimizer state. + + Expects the checkpoint at: + //.pk + """ + from collections import OrderedDict + + path_name = os.path.join( + pretrained_model_dir, pretrained_model_name, pretrained_model_name + ".pk" + ) + if not os.path.isfile(path_name): + raise FileNotFoundError( + f"Pretrained checkpoint not found at '{path_name}'. " + "Verify --pretrained_model_dir and --pretrained_model_name." + ) + + map_location = {"cuda:%d" % 0: str(get_device())} + info(f"[FT] Loading pretrained weights from: {path_name}") + checkpoint = torch.load(path_name, map_location=map_location) + state_dict = checkpoint["model_state_dict"] + + # Unwrap DDP prefix if needed + target = model.module if hasattr(model, "module") else model + own_keys = set(target.state_dict().keys()) + if own_keys and not next(iter(own_keys)).startswith("module"): + # Remove 'module.' prefix coming from a DDP-saved state dict + new_sd = OrderedDict() + for k, v in state_dict.items(): + new_sd[k.replace("module.", "", 1)] = v + state_dict = new_sd + + missing, unexpected = target.load_state_dict(state_dict, strict=False) + if missing: + info(f"[FT] WARNING: missing keys in checkpoint: {missing[:5]}{'...' if len(missing) > 5 else ''}") + if unexpected: + info(f"[FT] WARNING: unexpected keys in checkpoint: {unexpected[:5]}{'...' if len(unexpected) > 5 else ''}") + info("[FT] Pretrained weights loaded successfully.") + + +def detect_resume_state(log_name: str, logs_root: str = "./logs"): + """Detect whether a previous checkpoint exists and infer next epoch start.""" + run_dir = os.path.join(logs_root, log_name) + latest_link = os.path.join(run_dir, f"{log_name}.pk") + has_checkpoint = os.path.isfile(latest_link) + + epoch_ckpts = glob.glob(os.path.join(run_dir, f"{log_name}_epoch_*.pk")) + max_epoch = None + for ckpt in epoch_ckpts: + m = re.search(r"_epoch_(\d+)\.pk$", os.path.basename(ckpt)) + if m: + e = int(m.group(1)) + if max_epoch is None or e > max_epoch: + max_epoch = e + + if max_epoch is None: + epoch_start = 0 + else: + epoch_start = max_epoch + 1 + + return has_checkpoint, epoch_start, max_epoch + + +# --------------------------------------------------------------------------- +# Main entry point +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description="Fine-tune a pretrained HydraGNN OPF model.", + ) + # ── Dataset arguments ────────────────────────────────────────────────── + parser.add_argument( + "--inputfile", + type=str, + default="FT1_topology/config_HeteroSAGE_full.json", + help="Path to the fine-tuning JSON config (relative to this script's directory).", + ) + parser.add_argument( + "--modelname", + type=str, + default=None, + help="Log/checkpoint name for this fine-tuning run. Defaults to the config key.", + ) + parser.add_argument( + "--data_root", + type=str, + default="../dataset", + help="Root directory containing pre-serialised HDF5 data.", + ) + parser.add_argument( + "--node_target_type", + type=str, + default="bus", + choices=["bus", "generator"], + ) + # ── Data format ──────────────────────────────────────────────────────── + group = parser.add_mutually_exclusive_group() + group.add_argument("--hdf5", action="store_const", dest="format", const="hdf5") + group.add_argument("--pickle", action="store_const", dest="format", const="pickle") + parser.set_defaults(format="hdf5") + + # ── Training overrides ───────────────────────────────────────────────── + parser.add_argument("--num_epoch", type=int, default=None) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--learning_rate", type=float, default=None) + + # ── Fine-tuning specific ─────────────────────────────────────────────── + parser.add_argument( + "--pretrained_model_dir", + type=str, + default="../pretrained_models", + help="Directory containing pretrained model subdirectories.", + ) + parser.add_argument( + "--pretrained_model_name", + type=str, + required=True, + choices=["HeteroSAGE_best", "HeteroHEAT_best"], + help="Name of the pretrained model (must match a subdirectory in pretrained_model_dir).", + ) + parser.add_argument( + "--finetune_regime", + type=str, + default="full", + choices=["full", "partial", "head_only"], + help=( + "Freeze regime: " + "'full' trains all parameters, " + "'partial' trains last conv layer + head, " + "'head_only' trains only the prediction head." + ), + ) + parser.add_argument( + "--no_pretrained", + action="store_true", + default=False, + help=( + "Skip loading pretrained weights — train from random initialisation " + "(baseline comparison). Implies --finetune_regime full." + ), + ) + parser.add_argument( + "--max_train_samples", + type=int, + default=None, + help=( + "Truncate the training split to this many samples before training. " + "Val/test splits are unchanged. Used for data-efficiency sweeps." + ), + ) + parser.add_argument( + "--resume_if_exists", + action="store_true", + default=False, + help=( + "If logs//.pk exists, resume model+optimizer " + "from checkpoint and continue from the next epoch." + ), + ) + parser.add_argument( + "--ddp_find_unused_parameters", + dest="ddp_find_unused_parameters", + action="store_true", + help=( + "Enable DDP unused-parameter detection. This is safer for models " + "with conditional branches where not all parameters participate in " + "every step." + ), + ) + parser.add_argument( + "--no_ddp_find_unused_parameters", + dest="ddp_find_unused_parameters", + action="store_false", + help="Disable DDP unused-parameter detection.", + ) + parser.set_defaults(ddp_find_unused_parameters=True) + + args = parser.parse_args() + if args.no_pretrained: + args.finetune_regime = "full" + + # ── Resolve paths ────────────────────────────────────────────────────── + script_dir = os.path.dirname(os.path.abspath(__file__)) + input_filename = os.path.join(script_dir, args.inputfile) + pretrained_model_dir = os.path.join(script_dir, args.pretrained_model_dir) + data_root = os.path.join(script_dir, args.data_root) + + with open(input_filename) as f: + config = json.load(f) + + # ── Apply CLI overrides ──────────────────────────────────────────────── + if args.num_epoch is not None: + config["NeuralNetwork"]["Training"]["num_epoch"] = args.num_epoch + if args.batch_size is not None: + config["NeuralNetwork"]["Training"]["batch_size"] = args.batch_size + if args.learning_rate is not None: + config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"] = args.learning_rate + + arch_config = config["NeuralNetwork"]["Architecture"] + + edge_dim = _resolve_edge_dim(config) + + # ── Distributed initialisation ───────────────────────────────────────── + hydragnn.utils.distributed.setup_ddp() + + # ── Log name ─────────────────────────────────────────────────────────── + ft_tag = f"{args.pretrained_model_name}_{args.finetune_regime}" + if args.modelname is not None: + log_name = args.modelname + else: + log_name = f"finetune_{ft_tag}" + hydragnn.utils.print.setup_log(log_name) + + # Optional auto-resume from existing checkpoint for walltime-limited runs. + resume_exists = False + resume_epoch_start = 0 + resume_max_epoch = None + if args.resume_if_exists: + resume_exists, resume_epoch_start, resume_max_epoch = detect_resume_state(log_name) + if resume_exists: + config["NeuralNetwork"]["Training"]["epoch_start"] = resume_epoch_start + if resume_max_epoch is None: + info(f"[FT] Found existing checkpoint for '{log_name}', resuming from epoch_start=0") + else: + info( + f"[FT] Found existing checkpoint for '{log_name}' at epoch {resume_max_epoch}; " + f"resuming from epoch_start={resume_epoch_start}" + ) + else: + info(f"[FT] No existing checkpoint found for '{log_name}'; starting fresh") + + # ── Load serialised datasets ─────────────────────────────────────────── + # The dataset is expected to have been serialised by the preprocessing + # step of train_opf_solution_heterogeneous.py --preonly --hdf5. + # The modelname used during preprocessing is stored in the config as + # 'ft_data_modelname' or falls back to 'OPF_Solution_Hetero'. + data_modelname = config.get("ft_data_modelname", "OPF_Solution_Hetero") + basedir = os.path.join(data_root, f"{data_modelname}.h5") + if not os.path.isdir(basedir): + raise RuntimeError( + f"Pre-serialised HDF5 dataset not found at '{basedir}'. " + "Run the preprocessing step first:\n" + " python ../train_opf_solution_heterogeneous.py " + "--preonly --hdf5 --case_name [--max_samples N]" + ) + trainset = HDF5Dataset(basedir, "trainset") + valset = HDF5Dataset(basedir, "valset") + testset = HDF5Dataset(basedir, "testset") + + # ── Adapt datasets ───────────────────────────────────────────────────── + resolved_node_target_type = _resolve_node_target_type( + trainset[0], args.node_target_type + ) + args.node_target_type = resolved_node_target_type + config["NeuralNetwork"]["Architecture"]["node_target_type"] = args.node_target_type + validate_voi_node_features(config, args.node_target_type) + + trainset = EdgeAttrDatasetAdapter(trainset, edge_dim=edge_dim) + valset = EdgeAttrDatasetAdapter(valset, edge_dim=edge_dim) + testset = EdgeAttrDatasetAdapter(testset, edge_dim=edge_dim) + + trainset = NodeTargetDatasetAdapter(trainset, args.node_target_type, edge_dim=edge_dim) + valset = NodeTargetDatasetAdapter(valset, args.node_target_type, edge_dim=edge_dim) + testset = NodeTargetDatasetAdapter(testset, args.node_target_type, edge_dim=edge_dim) + + # Optionally truncate trainset for data-efficiency sweep + if args.max_train_samples is not None and len(trainset) > args.max_train_samples: + from torch.utils.data import Subset + + trainset = Subset(trainset, list(range(args.max_train_samples))) + info( + f"Truncated trainset to {args.max_train_samples} samples " + "for data-efficiency sweep." + ) + + info( + "trainset / valset / testset sizes: %d / %d / %d" + % (len(trainset), len(valset), len(testset)) + ) + + (train_loader, val_loader, test_loader) = hydragnn.preprocess.create_dataloaders( + trainset, valset, testset, + config["NeuralNetwork"]["Training"]["batch_size"], + ) + train_loader = NodeBatchAdapter(train_loader, args.node_target_type, edge_dim=edge_dim) + val_loader = NodeBatchAdapter(val_loader, args.node_target_type, edge_dim=edge_dim) + test_loader = NodeBatchAdapter(test_loader, args.node_target_type, edge_dim=edge_dim) + + config = update_config(config, train_loader, val_loader, test_loader) + arch_config = config["NeuralNetwork"]["Architecture"] + + if arch_config.get("mpnn_type") == "HeteroPNA" and not arch_config.get("pna_deg"): + pna_deg = compute_pna_deg_for_hetero_dataset(trainset, verbosity=2) + arch_config["pna_deg"] = pna_deg + arch_config["max_neighbours"] = max(0, len(pna_deg) - 1) + + config = _to_jsonable(config) + save_config(config, log_name) + + # ── Create model ─────────────────────────────────────────────────────── + node_input_dims = arch_config.get("node_input_dims") + if node_input_dims is None: + raise RuntimeError("Missing NeuralNetwork.Architecture.node_input_dims in config.") + + try: + metadata = trainset[0].metadata() + except Exception: + metadata = None + + model = hydragnn.models.create_model_config( + config=config["NeuralNetwork"], + verbosity=config["Verbosity"]["level"], + metadata=metadata, + node_input_dims=node_input_dims, + ) + + # ── Optionally wrap with domain-loss ─────────────────────────────────── + domain_loss_config = config["NeuralNetwork"]["Training"].get("DomainLoss") + if domain_loss_config is not None and domain_loss_config.get("enabled", False): + model = OPFEnhancedModelWrapper( + model, + OPFDomainLoss(domain_loss_config, node_target_type=args.node_target_type), + ) + + # ── Load pretrained weights (before freezing, before optimizer) ──────── + if not args.no_pretrained: + load_pretrained_weights(model, args.pretrained_model_name, pretrained_model_dir) + else: + info("[FT] --no_pretrained: starting from random initialisation (baseline).") + + # ── Apply freeze regime ──────────────────────────────────────────────── + apply_freeze_regime(model, args.finetune_regime) + + # ── Create optimizer over trainable parameters only ──────────────────── + trainable_params = [p for p in model.parameters() if p.requires_grad] + n_total = sum(p.numel() for p in model.parameters()) + n_trainable = sum(p.numel() for p in trainable_params) + info( + f"[FT] Regime '{args.finetune_regime}': " + f"{n_trainable:,} / {n_total:,} parameters trainable " + f"({100.0 * n_trainable / max(n_total, 1):.1f}%)" + ) + + learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"] + optimizer = torch.optim.AdamW(trainable_params, lr=learning_rate) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=5, min_lr=1e-6 + ) + + # ── Wrap in DDP ──────────────────────────────────────────────────────── + # Some FT3 configurations can leave a subset of parameters unused on a + # given step, which can trigger DDP reduction errors when this flag is off. + info( + "[FT] DDP find_unused_parameters=" + f"{args.ddp_find_unused_parameters}" + ) + model, optimizer = hydragnn.utils.distributed.distributed_model_wrapper( + model, optimizer, + config["Verbosity"]["level"], + find_unused_parameters=args.ddp_find_unused_parameters, + ) + + if args.resume_if_exists and resume_exists: + try: + load_existing_model(model, log_name, path="./logs/", optimizer=optimizer) + info(f"[FT] Resume loaded from logs/{log_name}/{log_name}.pk") + except Exception as exc: + # Checkpoints can become incompatible after architecture/config changes. + # Fall back to fresh training for this run instead of exiting early. + config["NeuralNetwork"]["Training"]["epoch_start"] = 0 + info( + "[FT] Resume checkpoint exists but failed to load; " + f"continuing fresh for this run. Error: {exc}" + ) + + print_model(model) + + # ── TensorBoard + CSV training-curve writer ─────────────────────────── + _tb_writer = model_utils.get_summary_writer(log_name) + _csv_path = os.path.join("logs", log_name, "training_curve.csv") + writer = EpochCSVWriter(_tb_writer, _csv_path) + + precision = config["NeuralNetwork"]["Training"].get("precision", "fp32") + + # ── Train ────────────────────────────────────────────────────────────── + hydragnn.train.train_validate_test( + model, + optimizer, + train_loader, + val_loader, + test_loader, + writer, + scheduler, + config["NeuralNetwork"], + log_name, + config["Verbosity"]["level"], + create_plots=False, + precision=precision, + ) + + # ── Flush domain-loss log ────────────────────────────────────────────── + _inner = model.module if hasattr(model, "module") else model + if isinstance(_inner, OPFEnhancedModelWrapper): + _inner._flush_epoch_log(_inner._last_seen_epoch) + + # ── Save final checkpoint ────────────────────────────────────────────── + model_utils.save_model(model, optimizer, log_name) + hydragnn.utils.profiling_and_tracing.print_timers(config["Verbosity"]["level"]) + + # ── Post-training regression evaluation ─────────────────────────────── + device = get_device() + _cfg_voi = config["NeuralNetwork"]["Variables_of_interest"] + _output_names = _cfg_voi.get("output_names", None) + if _output_names and len(_output_names) == 1: + _out_dim = arch_config.get("output_dim", [2])[0] + if args.node_target_type == "bus": + _out_names = ["Va", "Vm"][:_out_dim] + else: + _out_names = ["Pg", "Qg"][:_out_dim] + else: + _out_names = _output_names or ["dim_0", "dim_1"] + + info(f"[FT] Post-training regression eval on test set (dims: {_out_names})...") + test_metrics = evaluate_ft3(model, test_loader, device, comm, + output_names=_out_names) + + if rank == 0 and test_metrics is not None: + _mse_strs = ", ".join( + f"{n}_MSE={test_metrics.get(n + '_mse', float('nan')):.5f}" + for n in _out_names + ) + info(f"[FT] Test metrics — {_mse_strs} overall_MSE={test_metrics['overall_mse']:.5f}") + run_meta = { + "ft_strategy": config.get("_ft_strategy", "unknown"), + "arch": arch_config["mpnn_type"], + "regime": args.finetune_regime, + "pretrained": not args.no_pretrained, + "pretrained_model": args.pretrained_model_name if not args.no_pretrained else "none", + "node_target_type": args.node_target_type, + "num_epoch": config["NeuralNetwork"]["Training"]["num_epoch"], + "learning_rate": config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"], + "config_file": args.inputfile, + } + save_run_results(log_name, run_meta, test_metrics) + + comm.Barrier() + + writer.close() + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/examples/opf/finetune/train_opf_ft1_classify.py b/examples/opf/finetune/train_opf_ft1_classify.py new file mode 100644 index 000000000..99fbc7447 --- /dev/null +++ b/examples/opf/finetune/train_opf_ft1_classify.py @@ -0,0 +1,449 @@ +"""Fine-tuning script for FT1: OPF feasibility binary classification. + +Loads a pretrained HeteroSAGE or HeteroHEAT checkpoint and fine-tunes it for +graph-level binary feasibility prediction. + +The dataset must first be generated with generate_infeasible_samples.py, which +produces a balanced HDF5 dataset where: + - feasible samples: data.y = [1.0] (original OPF solutions) + - infeasible samples: data.y = [0.0] (load features scaled by overload_factor) + +The model uses a graph-level output head (pool all node embeddings → MLP → 1 +logit) with binary cross-entropy with logits loss. Inference: sigmoid(logit) +> 0.5 → feasible. + +Example:: + + python train_opf_ft1_classify.py \\ + --inputfile FT1_feasibility_classification/config_HeteroSAGE_full.json \\ + --data_root ../dataset \\ + --pretrained_model_dir ../pretrained_models \\ + --pretrained_model_name HeteroSAGE_best \\ + --finetune_regime full \\ + --num_epoch 50 \\ + --learning_rate 1e-4 +""" + +import os +import sys +import json +import argparse + +# Make examples/opf importable +_OPF_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..") +sys.path.insert(0, _OPF_DIR) + +import torch +import torch.distributed as dist +from collections import OrderedDict +from mpi4py import MPI + +import hydragnn +import hydragnn.utils.model as model_utils +from hydragnn.utils.model import print_model +from hydragnn.utils.distributed import get_device +from hydragnn.utils.input_config_parsing.config_utils import update_config +from hydragnn.utils.input_config_parsing import save_config +from hydragnn.utils.datasets.hdf5dataset import HDF5Dataset + +from ft_utils import EpochCSVWriter, evaluate_ft1, save_run_results + +from opf_solution_utils import ( + EdgeAttrDatasetAdapter, + OPFEnhancedModelWrapper, + info, +) + + +# --------------------------------------------------------------------------- +# Utility helpers (duplicated from train_opf_finetune.py for independence) +# --------------------------------------------------------------------------- + +def _to_jsonable(obj): + import numpy as np + if isinstance(obj, torch.Tensor): + return obj.item() if obj.numel() == 1 else obj.tolist() + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, np.generic): + return obj.item() + if isinstance(obj, dict): + return {k: _to_jsonable(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_to_jsonable(v) for v in obj] + return obj + + +def _resolve_edge_dim(config): + arch = config.get("NeuralNetwork", {}).get("Architecture", {}) + return arch.get("edge_dim", {"ac_line": 9, "transformer": 11}) + + +def apply_freeze_regime(model, regime: str): + """Freeze model parameters according to the fine-tuning regime. + + Called after loading pretrained weights, before creating the optimizer. + + Parameters + ---------- + model : nn.Module (not yet DDP-wrapped) + regime : str — 'full', 'partial', or 'head_only' + """ + if regime == "full": + return + + inner = model + if isinstance(inner, OPFEnhancedModelWrapper): + inner = inner.model + + if regime == "head_only": + for name, param in inner.named_parameters(): + if "heads_NN" not in name and "graph_shared" not in name: + param.requires_grad_(False) + n = sum(p.numel() for p in inner.parameters() if p.requires_grad) + info(f"[FT1] head_only: {n:,} trainable parameters (graph head only)") + return + + if regime == "partial": + for name, param in inner.named_parameters(): + if "input_projectors" in name or "node_embedders" in name: + param.requires_grad_(False) + n_conv = len(inner.graph_convs) + for i, (conv, feat) in enumerate(zip(inner.graph_convs, inner.feature_layers)): + if i < n_conv - 1: + for p in conv.parameters(): + p.requires_grad_(False) + for p in feat.parameters(): + p.requires_grad_(False) + n = sum(p.numel() for p in inner.parameters() if p.requires_grad) + info( + f"[FT1] partial: {n:,} trainable parameters " + f"(last conv layer + graph head; {n_conv-1} of {n_conv} conv layers frozen)" + ) + return + + raise ValueError( + f"Unknown finetune_regime '{regime}'. Choose from: full, partial, head_only." + ) + + +def load_pretrained_weights(model, pretrained_model_name: str, pretrained_model_dir: str): + """Load conv-layer weights from a pretrained node-level regression checkpoint. + + The prediction head is NOT loaded (strict=False), so the pretrained encoder + is reused while the graph-level classification head starts from scratch. + + Checkpoint expected at: + //.pk + """ + path = os.path.join( + pretrained_model_dir, + pretrained_model_name, + pretrained_model_name + ".pk", + ) + if not os.path.isfile(path): + raise FileNotFoundError( + f"Pretrained checkpoint not found: '{path}'. " + "Check --pretrained_model_dir and --pretrained_model_name." + ) + + map_loc = {"cuda:%d" % 0: str(get_device())} + info(f"[FT1] Loading pretrained weights from: {path}") + ckpt = torch.load(path, map_location=map_loc) + state_dict = ckpt["model_state_dict"] + + target = model.module if hasattr(model, "module") else model + own_keys = set(target.state_dict().keys()) + if own_keys and not next(iter(own_keys)).startswith("module"): + new_sd = OrderedDict() + for k, v in state_dict.items(): + new_sd[k.replace("module.", "", 1)] = v + state_dict = new_sd + + missing, unexpected = target.load_state_dict(state_dict, strict=False) + # Graph head keys will be missing (new head); report for transparency + head_keys = [k for k in missing if "heads_NN" in k or "graph_shared" in k] + other_miss = [k for k in missing if k not in head_keys] + info( + f"[FT1] Loaded weights: {len(missing)} missing keys " + f"({len(head_keys)} graph-head keys expected), " + f"{len(unexpected)} unexpected keys." + ) + if other_miss: + info(f"[FT1] Non-head missing keys: {other_miss[:5]}" + f"{'...' if len(other_miss) > 5 else ''}") + info("[FT1] Pretrained encoder loaded successfully.") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + comm = MPI.COMM_WORLD + rank = comm.Get_rank() + world_size = comm.Get_size() + + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description="FT1: fine-tune pretrained HydraGNN OPF model for feasibility classification.", + ) + parser.add_argument( + "--inputfile", + type=str, + default="FT1_feasibility_classification/config_HeteroSAGE_full.json", + help="Fine-tuning config JSON (relative to this script's directory).", + ) + parser.add_argument( + "--modelname", + type=str, + default=None, + help="Log/checkpoint name. Defaults to a derived tag.", + ) + parser.add_argument( + "--data_root", + type=str, + default="../dataset", + help="Root directory containing the FT1 HDF5 dataset.", + ) + parser.add_argument("--num_epoch", type=int, default=None) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--learning_rate", type=float, default=None) + parser.add_argument( + "--pretrained_model_dir", + type=str, + default="../pretrained_models", + ) + parser.add_argument( + "--pretrained_model_name", + type=str, + required=True, + choices=["HeteroSAGE_best", "HeteroHEAT_best"], + ) + parser.add_argument( + "--finetune_regime", + type=str, + default="full", + choices=["full", "partial", "head_only"], + ) + parser.add_argument( + "--no_pretrained", + action="store_true", + default=False, + help=( + "Skip loading pretrained weights — train from random initialisation " + "(baseline comparison). Implies --finetune_regime full." + ), + ) + parser.add_argument( + "--max_train_samples", + type=int, + default=None, + help=( + "Truncate the training split to this many samples before training. " + "Val/test splits are unchanged. Used for data-efficiency sweeps." + ), + ) + + args = parser.parse_args() + if args.no_pretrained: + args.finetune_regime = "full" # no freezing for baseline + + # ── Resolve paths ────────────────────────────────────────────────────── + script_dir = os.path.dirname(os.path.abspath(__file__)) + input_filename = os.path.join(script_dir, args.inputfile) + pretrained_model_dir = os.path.join(script_dir, args.pretrained_model_dir) + data_root = os.path.join(script_dir, args.data_root) + + with open(input_filename) as f: + config = json.load(f) + + # ── CLI overrides ────────────────────────────────────────────────────── + if args.num_epoch is not None: + config["NeuralNetwork"]["Training"]["num_epoch"] = args.num_epoch + if args.batch_size is not None: + config["NeuralNetwork"]["Training"]["batch_size"] = args.batch_size + if args.learning_rate is not None: + config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"] = args.learning_rate + + arch_config = config["NeuralNetwork"]["Architecture"] + edge_dim = _resolve_edge_dim(config) + + # ── Distributed init ─────────────────────────────────────────────────── + hydragnn.utils.distributed.setup_ddp() + + # ── Log name ─────────────────────────────────────────────────────────── + ft_tag = f"{args.pretrained_model_name}_{args.finetune_regime}" + log_name = args.modelname if args.modelname is not None else f"FT1_classify_{ft_tag}" + hydragnn.utils.print.setup_log(log_name) + + # ── Load HDF5 dataset ────────────────────────────────────────────────── + data_modelname = config.get("ft_data_modelname", "FT1_feasibility_data") + basedir = os.path.join(data_root, f"{data_modelname}.h5") + if not os.path.isdir(basedir): + raise RuntimeError( + f"FT1 dataset not found at '{basedir}'.\n" + "Generate it first with:\n" + " python generate_infeasible_samples.py " + "--src_dir ../dataset/.h5 " + "--out_dir ../dataset/FT1_feasibility_data.h5" + ) + + trainset = HDF5Dataset(basedir, "trainset") + valset = HDF5Dataset(basedir, "valset") + testset = HDF5Dataset(basedir, "testset") + + # Validate edge feature shapes (no node-level target injection needed — + # data.y is already the graph-level label set by generate_infeasible_samples.py) + trainset = EdgeAttrDatasetAdapter(trainset, edge_dim=edge_dim) + valset = EdgeAttrDatasetAdapter(valset, edge_dim=edge_dim) + testset = EdgeAttrDatasetAdapter(testset, edge_dim=edge_dim) + + # Optionally truncate trainset for data-efficiency sweep + if args.max_train_samples is not None and len(trainset) > args.max_train_samples: + from torch.utils.data import Subset + + trainset = Subset(trainset, list(range(args.max_train_samples))) + info( + f"Truncated trainset to {args.max_train_samples} samples " + "for data-efficiency sweep." + ) + + info( + "FT1 dataset sizes: train=%d val=%d test=%d" + % (len(trainset), len(valset), len(testset)) + ) + sample0 = trainset[0] + info( + f" Sample node types: {sample0.node_types} | " + f"y = {sample0.y.tolist()}" + ) + + # ── DataLoaders ──────────────────────────────────────────────────────── + (train_loader, val_loader, test_loader) = hydragnn.preprocess.create_dataloaders( + trainset, valset, testset, + config["NeuralNetwork"]["Training"]["batch_size"], + ) + + # ── Update config from data ──────────────────────────────────────────── + # update_config handles graph-level outputs without y_loc: it reads + # output_dim directly from config["Variables_of_interest"]["output_dim"]. + config = update_config(config, train_loader, val_loader, test_loader) + arch_config = config["NeuralNetwork"]["Architecture"] + + config = _to_jsonable(config) + save_config(config, log_name) + + # ── Create model ─────────────────────────────────────────────────────── + node_input_dims = arch_config.get("node_input_dims") + if node_input_dims is None: + raise RuntimeError( + "node_input_dims not found in config after update_config. " + "Ensure the dataset contains HeteroData with node-type features." + ) + + try: + metadata = trainset[0].metadata() + except Exception: + metadata = None + + model = hydragnn.models.create_model_config( + config=config["NeuralNetwork"], + verbosity=config["Verbosity"]["level"], + metadata=metadata, + node_input_dims=node_input_dims, + ) + + # ── Optionally load pretrained encoder weights ──────────────────────── + if not args.no_pretrained: + load_pretrained_weights(model, args.pretrained_model_name, pretrained_model_dir) + else: + info("[FT1] --no_pretrained: starting from random initialisation (baseline).") + + # ── Apply freeze regime ──────────────────────────────────────────────── + apply_freeze_regime(model, args.finetune_regime) + + # ── Optimizer ───────────────────────────────────────────────────────── + trainable_params = [p for p in model.parameters() if p.requires_grad] + n_total = sum(p.numel() for p in model.parameters()) + n_trainable = sum(p.numel() for p in trainable_params) + info( + f"[FT1] Regime '{args.finetune_regime}': " + f"{n_trainable:,} / {n_total:,} parameters trainable " + f"({100.0 * n_trainable / max(n_total, 1):.1f}%)" + ) + + lr = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"] + optimizer = torch.optim.AdamW(trainable_params, lr=lr) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=5, min_lr=1e-6 + ) + + # ── DDP wrap ─────────────────────────────────────────────────────────── + model, optimizer = hydragnn.utils.distributed.distributed_model_wrapper( + model, optimizer, + config["Verbosity"]["level"], + find_unused_parameters=(args.finetune_regime != "full"), + ) + + print_model(model) + + # ── TensorBoard + CSV training-curve writer ─────────────────────────── + _tb_writer = model_utils.get_summary_writer(log_name) + _csv_path = os.path.join("logs", log_name, "training_curve.csv") + writer = EpochCSVWriter(_tb_writer, _csv_path) + + precision = config["NeuralNetwork"]["Training"].get("precision", "fp32") + + # ── Train ────────────────────────────────────────────────────────────── + hydragnn.train.train_validate_test( + model, + optimizer, + train_loader, + val_loader, + test_loader, + writer, + scheduler, + config["NeuralNetwork"], + log_name, + config["Verbosity"]["level"], + create_plots=False, + precision=precision, + ) + + # ── Save checkpoint ──────────────────────────────────────────────────── + model_utils.save_model(model, optimizer, log_name) + hydragnn.utils.profiling_and_tracing.print_timers(config["Verbosity"]["level"]) + + # ── Post-training evaluation ─────────────────────────────────────────── + device = get_device() + info("[FT1] Running post-training classification evaluation on test set...") + test_metrics = evaluate_ft1(model, test_loader, device, comm) + + if rank == 0 and test_metrics is not None: + info( + f"[FT1] Test metrics — " + f"BCE={test_metrics['bce']:.4f} " + f"Acc={test_metrics['accuracy']:.4f} " + f"F1={test_metrics['f1']:.4f} " + f"AUC={test_metrics['auc_roc']}" + ) + run_meta = { + "ft_strategy": "FT1_feasibility_classification", + "arch": arch_config["mpnn_type"], + "regime": args.finetune_regime, + "pretrained": not args.no_pretrained, + "pretrained_model": args.pretrained_model_name if not args.no_pretrained else "none", + "num_epoch": config["NeuralNetwork"]["Training"]["num_epoch"], + "learning_rate": config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"], + "config_file": args.inputfile, + } + save_run_results(log_name, run_meta, test_metrics) + + comm.Barrier() + + writer.close() + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/examples/opf/finetune/watch_and_chain_bundles.sh b/examples/opf/finetune/watch_and_chain_bundles.sh new file mode 100755 index 000000000..2060cd5a2 --- /dev/null +++ b/examples/opf/finetune/watch_and_chain_bundles.sh @@ -0,0 +1,152 @@ +#!/bin/bash +# ============================================================================= +# Bundled per-(arch, method, bundle) submission for OPF FT1 + FT3. +# +# Each Slurm job runs ALL sample sizes in a "bundle" inside one allocation. +# Bundles are sized so that the total expected wall-time fits the selected +# QOS walltime cap (rough estimate: ~0.5 s per train sample). +# +# Per-cell isolation is preserved: the underlying job scripts already write +# each (arch, method, N) into its own logs/{LOG_NAME}/ directory, so cells +# cannot overwrite each other. +# +# If a bundle is killed mid-way, cells whose results.json already wrote will +# be skipped on the next run; the existing per-N watcher (or this same +# bundle watcher) will pick up the unfinished cells automatically. +# ============================================================================= +set -u + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune + +ACCOUNT=eng164 +QOS=debug +PARTITION=batch +TIME=02:00:00 +NODES=1 + +ARCHES=(HeteroSAGE HeteroHEAT) +METHODS=(full partial head_only scratch) + +# Bundles, ordered cheapest → most expensive within each group. +# A: 5 cheap N's, ~120 min worst case +# B: two medium N's, ~145 min +# C: 50000 alone, ~90 min +# D: 100000 alone, ~150 min +# E: 270000 alone (FT3 only) +FT1_BUNDLES=( + "100 500 1000 2500 5000" + "10000 25000" + "50000" + "100000" +) +FT3_BUNDLES=( + "100 500 1000 2500 5000" + "10000 25000" + "50000" + "100000" + "270000" +) + +LOG=$FT_DIR/ft_bundles_chain.log +: > "$LOG" + +# Compute LOG_NAME for one (task, arch, method, n) — must match the rules +# baked into the underlying job scripts. +log_name_for() { + local task=$1 arch=$2 method=$3 n=$4 + local scratch_tag="" regime="$method" + if [[ "$method" == "scratch" ]]; then + scratch_tag="_scratch" + regime="full" + fi + if [[ "$task" == "FT1" ]]; then + echo "FT1_feasibility_${arch}_${regime}${scratch_tag}_n${n}" + else + echo "finetune_FT3_contingency_${arch}_${regime}${scratch_tag}_n${n}" + fi +} + +# Filter out N's that already have a results.json on disk. +remaining_ns() { + local task=$1 arch=$2 method=$3 + shift 3 + local out="" n logname + for n in "$@"; do + logname=$(log_name_for "$task" "$arch" "$method" "$n") + if [[ ! -f "$FT_DIR/logs/$logname/results.json" ]]; then + out+="$n " + fi + done + echo "$out" | sed 's/[[:space:]]*$//' +} + +submit_bundle() { + local task=$1 arch=$2 method=$3 + shift 3 + local ns=("$@") + local sizes_str="${ns[*]}" + + local script jobname + case "$task" in + FT1) script=$FT_DIR/job-frontier-FT1-single-method.sh ;; + FT3) script=$FT_DIR/job-frontier-FT3-single-method.sh ;; + *) echo "Unknown task $task" >> "$LOG"; return 1 ;; + esac + + local first="${ns[0]}" last="${ns[-1]}" + jobname="${task}-${arch}-${method}-bundle-${first}-${last}" + local out=$HYDRAGNN_ROOT/${jobname}-%j.out + + # Wait for an available slot in the selected QOS. + while squeue -u "$USER" -h --qos=$QOS -o "%i" 2>/dev/null | grep -q .; do + echo "[$(date +%H:%M:%S)] $QOS QOS busy — waiting 30s before $jobname" >> "$LOG" + sleep 30 + done + + local jid + jid=$(sbatch --parsable \ + --account=$ACCOUNT --partition=$PARTITION --qos=$QOS \ + --job-name="$jobname" \ + --nodes=$NODES --time=$TIME \ + --output="$out" --error="$out" \ + --export=ALL,FT_ARCH=$arch,FT_METHOD=$method,SAMPLE_SIZES="$sizes_str",N_PER_RUN=1 \ + "$script" 2>>"$LOG") + if [[ -z "$jid" ]]; then + echo "[$(date +%H:%M:%S)] [submit-FAILED] $jobname N=[$sizes_str]" >> "$LOG" + return 1 + fi + echo "[$(date +%H:%M:%S)] [submit] $jobname → job $jid N=[$sizes_str]" >> "$LOG" + + # Poll until the job leaves the queue. + while squeue -j "$jid" -h -o "%T" 2>/dev/null | grep -q .; do + sleep 30 + done + echo "[$(date +%H:%M:%S)] $jobname (job $jid) finished." >> "$LOG" +} + +run_task_bundles() { + local task=$1 arch=$2 method=$3 + shift 3 + local bundles=("$@") + local b ns + for b in "${bundles[@]}"; do + # Filter already-done cells. + ns=$(remaining_ns "$task" "$arch" "$method" $b) + if [[ -z "$ns" ]]; then + echo "[$(date +%H:%M:%S)] [skip-bundle] $task $arch $method [$b] (all cells done)" >> "$LOG" + continue + fi + # shellcheck disable=SC2086 + submit_bundle "$task" "$arch" "$method" $ns + done +} + +for ARCH in "${ARCHES[@]}"; do + for METHOD in "${METHODS[@]}"; do + run_task_bundles FT1 "$ARCH" "$METHOD" "${FT1_BUNDLES[@]}" + run_task_bundles FT3 "$ARCH" "$METHOD" "${FT3_BUNDLES[@]}" + done +done + +echo "[$(date +%H:%M:%S)] All bundles submitted and finished." >> "$LOG" diff --git a/examples/opf/finetune/watch_and_chain_debug.sh b/examples/opf/finetune/watch_and_chain_debug.sh new file mode 100755 index 000000000..072b7328a --- /dev/null +++ b/examples/opf/finetune/watch_and_chain_debug.sh @@ -0,0 +1,63 @@ +#!/bin/bash +# ============================================================================= +# watch_and_chain_debug.sh +# +# Submits 16 FT single-method jobs to the debug QOS one at a time. +# Waits for each job to leave the queue (RUNNING→finished) before +# submitting the next one. Polls every 30 seconds. +# +# Usage: +# nohup bash watch_and_chain_debug.sh > /tmp/ft_debug_chain.log 2>&1 & +# tail -f /tmp/ft_debug_chain.log +# ============================================================================= +set -uo pipefail + +FT_DIR="$(cd "$(dirname "$0")" && pwd)" +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN +POLL_INTERVAL=30 + +log() { echo "[$(date '+%H:%M:%S')] $*"; } + +wait_for_job() { + local JID="$1" + while squeue -j "$JID" -h 2>/dev/null | grep -q "$JID"; do + log " Job $JID still in queue — waiting ${POLL_INTERVAL}s..." + sleep $POLL_INTERVAL + done + log " Job $JID finished." +} + +submit_debug() { + local TASK="$1" ARCH="$2" METHOD="$3" + local LABEL="${TASK}-${ARCH}-${METHOD}" + local SCRIPT="${FT_DIR}/job-frontier-${TASK}-single-method.sh" + + # Wait until debug QOS is free (no jobs of ours there) + while squeue -u "$USER" --qos=debug -h 2>/dev/null | grep -q "$USER"; do + log " Debug QOS busy — waiting ${POLL_INTERVAL}s before submitting $LABEL ..." + sleep $POLL_INTERVAL + done + + local ID + ID=$(sbatch \ + --account=eng164 --partition=batch --qos=debug \ + --job-name="$LABEL" \ + --nodes=1 --time=02:00:00 \ + --output="${HYDRAGNN_ROOT}/${LABEL}-%j.out" \ + --error="${HYDRAGNN_ROOT}/${LABEL}-%j.out" \ + --export=ALL,FT_ARCH=${ARCH},FT_METHOD=${METHOD},N_PER_RUN=1 \ + "$SCRIPT" | awk '{print $NF}') + log "[submit] $LABEL → job $ID" + wait_for_job "$ID" +} + +log "=== FT debug chain starting — 16 jobs ===" + +for ARCH in HeteroSAGE HeteroHEAT; do + for METHOD in full partial head_only scratch; do + submit_debug FT1 "$ARCH" "$METHOD" + submit_debug FT3 "$ARCH" "$METHOD" + done +done + +log "=== All 16 FT jobs done! ===" diff --git a/examples/opf/finetune/watch_and_chain_per_n.sh b/examples/opf/finetune/watch_and_chain_per_n.sh new file mode 100755 index 000000000..d2bd27709 --- /dev/null +++ b/examples/opf/finetune/watch_and_chain_per_n.sh @@ -0,0 +1,91 @@ +#!/bin/bash +# ============================================================================= +# Per-(arch, method, N) chained submission for OPF FT1 + FT3. +# Submits ONE 1-node job at a time, polls for completion, then submits +# the next. Each job runs a single sample size for a single (arch, method) +# combination so that it fits comfortably inside the selected QOS walltime. +# ============================================================================= +set -u + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN +FT_DIR=$HYDRAGNN_ROOT/examples/opf/finetune + +ACCOUNT=eng164 +QOS=debug +PARTITION=batch +TIME=02:00:00 +NODES=1 + +ARCHES=(HeteroSAGE HeteroHEAT) +METHODS=(full partial head_only scratch) +FT1_NS=(100 500 1000 2500 5000 10000 25000 50000 100000) +FT3_NS=(100 500 1000 2500 5000 10000 25000 50000 100000 270000) + +LOG=/lustre/orion/lrn078/proj-shared/HydraGNN/examples/opf/finetune/ft_per_n_chain.log +: > "$LOG" + +submit_one() { + local task=$1 arch=$2 method=$3 n=$4 + local script jobname + case "$task" in + FT1) script=$FT_DIR/job-frontier-FT1-single-method.sh ;; + FT3) script=$FT_DIR/job-frontier-FT3-single-method.sh ;; + *) echo "Unknown task $task" >> "$LOG"; return 1 ;; + esac + jobname="${task}-${arch}-${method}-n${n}" + local out=$HYDRAGNN_ROOT/${jobname}-%j.out + + # Skip if a results.json already exists for this exact (task, arch, method, n). + local scratch_tag="" + [[ "$method" == "scratch" ]] && scratch_tag="_scratch" + local regime="$method" + [[ "$method" == "scratch" ]] && regime="full" + local logname + if [[ "$task" == "FT1" ]]; then + logname="FT1_feasibility_${arch}_${regime}${scratch_tag}_n${n}" + else + logname="finetune_FT3_contingency_${arch}_${regime}${scratch_tag}_n${n}" + fi + if [[ -f "$FT_DIR/logs/$logname/results.json" ]]; then + echo "[$(date +%H:%M:%S)] [skip] $jobname (results.json present)" >> "$LOG" + return 0 + fi + + # Wait for selected QOS slot (only 1 submitted job allowed). + while squeue -u "$USER" -h --qos=$QOS -o "%i" 2>/dev/null | grep -q .; do + echo "[$(date +%H:%M:%S)] $QOS QOS busy — waiting 30s before $jobname" >> "$LOG" + sleep 30 + done + + local jid + jid=$(sbatch --parsable \ + --account=$ACCOUNT --partition=$PARTITION --qos=$QOS \ + --job-name="$jobname" \ + --nodes=$NODES --time=$TIME \ + --output="$out" --error="$out" \ + --export=ALL,FT_ARCH=$arch,FT_METHOD=$method,SAMPLE_SIZES=$n,N_PER_RUN=1 \ + "$script" 2>>"$LOG") + if [[ -z "$jid" ]]; then + echo "[$(date +%H:%M:%S)] [submit-FAILED] $jobname" >> "$LOG" + return 1 + fi + echo "[$(date +%H:%M:%S)] [submit] $jobname → job $jid" >> "$LOG" + + while squeue -j "$jid" -h -o "%T" 2>/dev/null | grep -q .; do + sleep 30 + done + echo "[$(date +%H:%M:%S)] $jobname (job $jid) finished." >> "$LOG" +} + +for ARCH in "${ARCHES[@]}"; do + for METHOD in "${METHODS[@]}"; do + for N in "${FT1_NS[@]}"; do + submit_one FT1 "$ARCH" "$METHOD" "$N" + done + for N in "${FT3_NS[@]}"; do + submit_one FT3 "$ARCH" "$METHOD" "$N" + done + done +done + +echo "[$(date +%H:%M:%S)] All jobs submitted and finished." >> "$LOG" diff --git a/examples/opf/hpo_validation_loss_curves.png b/examples/opf/hpo_validation_loss_curves.png new file mode 100644 index 000000000..4108cb2e5 Binary files /dev/null and b/examples/opf/hpo_validation_loss_curves.png differ diff --git a/examples/opf/infer_opf_solution_heterogeneous.py b/examples/opf/infer_opf_solution_heterogeneous.py new file mode 100644 index 000000000..f59e71d6f --- /dev/null +++ b/examples/opf/infer_opf_solution_heterogeneous.py @@ -0,0 +1,436 @@ +"""Run inference for OPF heterogeneous node-level model and generate parity plots.""" + +import argparse +import json +import os +from mpi4py import MPI + +import matplotlib as mpl + +mpl.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.distributed as dist + +import hydragnn +from hydragnn.postprocess.postprocess import output_denormalize +from hydragnn.train.train_validate_test import test +from hydragnn.utils.distributed import setup_ddp +from hydragnn.utils.input_config_parsing.config_utils import update_config +from hydragnn.utils.model import load_existing_model + +from hydragnn.utils.datasets.pickledataset import SimplePickleDataset +from hydragnn.utils.datasets.hdf5dataset import HDF5Dataset + +from opf_solution_utils import ( + HeteroFromHomogeneousDataset, + NodeBatchAdapter, + NodeTargetDatasetAdapter, + compute_pna_deg_for_hetero_dataset, + validate_voi_node_features, + info, + resolve_edge_feature_schema, + resolve_node_target_type, +) + +try: + from hydragnn.utils.datasets.adiosdataset import AdiosDataset +except ImportError: + AdiosDataset = None + + +def _plot_parity_per_dim( + true_values, + predicted_values, + output_name, + output_dim, + out_dir, + prefix="test", +): + true_arr = true_values.detach().cpu().numpy() + pred_arr = predicted_values.detach().cpu().numpy() + if true_arr.ndim == 1: + true_arr = true_arr.reshape(-1, 1) + if pred_arr.ndim == 1: + pred_arr = pred_arr.reshape(-1, 1) + + if output_dim is None: + output_dim = true_arr.shape[1] + + total = true_arr.shape[0] + if total % output_dim != 0: + output_dim = true_arr.shape[1] + + true_arr = true_arr.reshape(-1, output_dim) + pred_arr = pred_arr.reshape(-1, output_dim) + + for dim in range(output_dim): + t = true_arr[:, dim] + p = pred_arr[:, dim] + fig, ax = plt.subplots(figsize=(5, 5)) + ax.scatter(t, p, s=8, alpha=0.6) + minv = float(np.min([t.min(), p.min()])) + maxv = float(np.max([t.max(), p.max()])) + ax.plot([minv, maxv], [minv, maxv], "r--", linewidth=1) + ax.set_xlabel("True") + ax.set_ylabel("Predicted") + ax.set_title(f"{output_name} dim {dim}") + ax.set_aspect("equal", adjustable="box") + fig.tight_layout() + fname = f"{prefix}_parity_{output_name}_dim{dim}.png" + fig.savefig(os.path.join(out_dir, fname), dpi=300) + plt.close(fig) + + +def _compute_mae_per_quantity(true_values, predicted_values, output_name, output_dim): + true_arr = true_values.detach().cpu().numpy() + pred_arr = predicted_values.detach().cpu().numpy() + if true_arr.ndim == 1: + true_arr = true_arr.reshape(-1, 1) + if pred_arr.ndim == 1: + pred_arr = pred_arr.reshape(-1, 1) + + if output_dim is None: + output_dim = true_arr.shape[1] + + total = true_arr.shape[0] + if total % output_dim != 0: + output_dim = true_arr.shape[1] + + true_arr = true_arr.reshape(-1, output_dim) + pred_arr = pred_arr.reshape(-1, output_dim) + + abs_err = np.abs(pred_arr - true_arr) + mae_per_dim = abs_err.mean(axis=0) + return { + "quantity": output_name, + "mae_overall": float(abs_err.mean()), + "mae_per_dim": [float(v) for v in mae_per_dim], + } + + +def _compute_diagnostics_per_quantity( + true_values, predicted_values, output_name, output_dim +): + true_arr = true_values.detach().cpu().numpy() + pred_arr = predicted_values.detach().cpu().numpy() + if true_arr.ndim == 1: + true_arr = true_arr.reshape(-1, 1) + if pred_arr.ndim == 1: + pred_arr = pred_arr.reshape(-1, 1) + + if output_dim is None: + output_dim = true_arr.shape[1] + + total = true_arr.shape[0] + if total % output_dim != 0: + output_dim = true_arr.shape[1] + + true_arr = true_arr.reshape(-1, output_dim) + pred_arr = pred_arr.reshape(-1, output_dim) + + residual = pred_arr - true_arr + abs_err = np.abs(residual) + + bias_per_dim = residual.mean(axis=0) + p50_per_dim = np.percentile(abs_err, 50, axis=0) + p90_per_dim = np.percentile(abs_err, 90, axis=0) + p99_per_dim = np.percentile(abs_err, 99, axis=0) + + high_true_bias_per_dim = [] + for dim in range(output_dim): + threshold = float(np.percentile(true_arr[:, dim], 90)) + mask = true_arr[:, dim] >= threshold + if np.any(mask): + high_true_bias_per_dim.append(float(residual[mask, dim].mean())) + else: + high_true_bias_per_dim.append(0.0) + + return { + "quantity": output_name, + "bias_per_dim": [float(v) for v in bias_per_dim], + "abs_error_p50_per_dim": [float(v) for v in p50_per_dim], + "abs_error_p90_per_dim": [float(v) for v in p90_per_dim], + "abs_error_p99_per_dim": [float(v) for v in p99_per_dim], + "high_true_bias_per_dim": high_true_bias_per_dim, + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "--inputfile", type=str, default="opf_solution_heterogeneous.json" + ) + parser.add_argument("--data_root", type=str, default="dataset") + parser.add_argument("--modelname", type=str, default="OPF_Solution") + parser.add_argument( + "--node_target_type", + type=str, + default="bus", + choices=["bus", "generator"], + ) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument( + "--config_from_log", + action="store_true", + help="Load config from logs//config.json when available", + ) + parser.set_defaults(config_from_log=True) + + group = parser.add_mutually_exclusive_group() + group.add_argument("--adios", action="store_const", dest="format", const="adios") + group.add_argument("--pickle", action="store_const", dest="format", const="pickle") + group.add_argument("--hdf5", action="store_const", dest="format", const="hdf5") + parser.set_defaults(format="pickle") + + args = parser.parse_args() + + comm = MPI.COMM_WORLD + comm_size, rank = setup_ddp() + + dirpwd = os.path.dirname(os.path.abspath(__file__)) + input_filename = os.path.join(dirpwd, args.inputfile) + log_config_path = os.path.join("./logs", args.modelname, "config.json") + + if args.config_from_log and os.path.isfile(log_config_path): + if rank == 0: + info(f"Loading config from {log_config_path}") + with open(log_config_path, "r") as f: + config = json.load(f) + else: + if not os.path.isfile(input_filename): + raise FileNotFoundError(f"Missing config file: {input_filename}") + with open(input_filename, "r") as f: + config = json.load(f) + + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + raw_edge_dim = arch_config.get("edge_dim") + if isinstance(raw_edge_dim, dict): + edge_dim = {str(k): int(v) for k, v in raw_edge_dim.items()} + edge_feature_schema = None + elif raw_edge_dim is not None: + edge_dim = int(raw_edge_dim) + names = arch_config.get("edge_feature_names") + if names: + edge_feature_schema = resolve_edge_feature_schema(names, edge_dim) + else: + edge_feature_schema = None + else: + raise RuntimeError("edge_dim must be specified in config.") + arch_config["edge_dim"] = edge_dim + + if "node_target_type" in config.get("NeuralNetwork", {}).get("Architecture", {}): + args.node_target_type = config["NeuralNetwork"]["Architecture"][ + "node_target_type" + ] + validate_voi_node_features(config, args.node_target_type) + + if args.batch_size is not None: + config["NeuralNetwork"]["Training"]["batch_size"] = args.batch_size + + if args.format == "adios": + if AdiosDataset is None: + raise RuntimeError("adios2 is not available in this environment.") + fname = os.path.join(dirpwd, "dataset", f"{args.modelname}.bp") + if not os.path.exists(fname): + raise FileNotFoundError(f"Missing ADIOS dataset: {fname}") + train_base = AdiosDataset(fname, "trainset", comm, var_config=None) + val_base = AdiosDataset(fname, "valset", comm, var_config=None) + test_base = AdiosDataset(fname, "testset", comm, var_config=None) + trainset = HeteroFromHomogeneousDataset(train_base, edge_dim=edge_dim) + valset = HeteroFromHomogeneousDataset(val_base, edge_dim=edge_dim) + testset = HeteroFromHomogeneousDataset(test_base, edge_dim=edge_dim) + elif args.format == "hdf5": + basedir = os.path.join(dirpwd, "dataset", f"{args.modelname}.h5") + if not os.path.isdir(basedir): + raise FileNotFoundError(f"Missing HDF5 dataset dir: {basedir}") + trainset = HDF5Dataset(basedir, "trainset") + valset = HDF5Dataset(basedir, "valset") + testset = HDF5Dataset(basedir, "testset") + else: + basedir = os.path.join(dirpwd, "dataset", f"{args.modelname}.pickle") + if not os.path.isdir(basedir): + raise FileNotFoundError(f"Missing pickle dataset dir: {basedir}") + trainset = SimplePickleDataset( + basedir=basedir, label="trainset", var_config=None + ) + valset = SimplePickleDataset(basedir=basedir, label="valset", var_config=None) + testset = SimplePickleDataset(basedir=basedir, label="testset", var_config=None) + + resolved_node_target_type = resolve_node_target_type( + trainset[0], args.node_target_type + ) + if resolved_node_target_type != args.node_target_type: + info( + f"Resolved node_target_type '{args.node_target_type}' -> '{resolved_node_target_type}'" + ) + args.node_target_type = resolved_node_target_type + config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {})[ + "node_target_type" + ] = args.node_target_type + validate_voi_node_features(config, args.node_target_type) + + trainset = NodeTargetDatasetAdapter( + trainset, args.node_target_type, edge_dim=edge_dim + ) + valset = NodeTargetDatasetAdapter(valset, args.node_target_type, edge_dim=edge_dim) + testset = NodeTargetDatasetAdapter( + testset, args.node_target_type, edge_dim=edge_dim + ) + + (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders( + trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"] + ) + + train_loader = NodeBatchAdapter( + train_loader, args.node_target_type, edge_dim=edge_dim + ) + val_loader = NodeBatchAdapter(val_loader, args.node_target_type, edge_dim=edge_dim) + test_loader = NodeBatchAdapter( + test_loader, args.node_target_type, edge_dim=edge_dim + ) + + config = update_config(config, train_loader, val_loader, test_loader) + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + if arch_config.get("mpnn_type") == "HeteroPNA" and not arch_config.get("pna_deg"): + info("Computing pna_deg for HeteroPNA from inference dataset") + pna_deg = compute_pna_deg_for_hetero_dataset(trainset, verbosity=2) + arch_config["pna_deg"] = pna_deg + arch_config["max_neighbours"] = max(0, len(pna_deg) - 1) + + metadata = None + try: + metadata = trainset[0].metadata() + except Exception as exc: + if rank == 0: + info(f"Unable to fetch hetero metadata: {exc}") + node_input_dims = ( + config.get("NeuralNetwork", {}).get("Architecture", {}).get("node_input_dims") + ) + if node_input_dims is None: + raise RuntimeError( + "Missing NeuralNetwork.Architecture.node_input_dims in config. " + "Add node_input_dims to the config to initialize node embedders." + ) + + model = hydragnn.models.create_model_config( + config=config["NeuralNetwork"], + verbosity=config["Verbosity"]["level"], + metadata=metadata, + node_input_dims=node_input_dims, + ) + + model = hydragnn.utils.distributed.distributed_model_wrapper( + model, None, config["Verbosity"]["level"] + )[0] + + load_existing_model(model, args.modelname) + + num_tasks = model.module.num_heads + test_error, task_errors, true_values, predicted_values = test( + test_loader, + model, + config["Verbosity"]["level"], + num_tasks=num_tasks, + precision=config["NeuralNetwork"]["Training"].get("precision", "fp32"), + ) + + if config["NeuralNetwork"]["Variables_of_interest"].get("denormalize_output"): + true_values, predicted_values = output_denormalize( + config["NeuralNetwork"]["Variables_of_interest"]["y_minmax"], + true_values, + predicted_values, + ) + + if rank == 0: + out_dir = os.path.join("./logs", args.modelname) + os.makedirs(out_dir, exist_ok=True) + var_config = config["NeuralNetwork"]["Variables_of_interest"] + output_names = var_config.get("output_names", None) + output_dims = var_config.get("output_dim", None) + + mae_metrics = [] + diagnostics_metrics = [] + for ihead in range(num_tasks): + name = output_names[ihead] if output_names else f"head{ihead}" + dim = output_dims[ihead] if output_dims else None + mae_metrics.append( + _compute_mae_per_quantity( + true_values[ihead], + predicted_values[ihead], + name, + dim, + ) + ) + diagnostics_metrics.append( + _compute_diagnostics_per_quantity( + true_values[ihead], + predicted_values[ihead], + name, + dim, + ) + ) + + metrics = { + "test_error": float(test_error.detach().cpu().item()), + "task_errors": task_errors.detach().cpu().tolist(), + "mae": mae_metrics, + "diagnostics": diagnostics_metrics, + } + with open(os.path.join(out_dir, "test_metrics.json"), "w") as f: + json.dump(metrics, f, indent=2) + + # Print key metrics to stdout so they are visible in batch job logs. + print(f"Inference test_error: {metrics['test_error']}", flush=True) + print(f"Inference task_errors: {metrics['task_errors']}", flush=True) + + for mae_entry in mae_metrics: + info( + "Inference MAE %s: overall=%g per_dim=%s" + % ( + mae_entry["quantity"], + mae_entry["mae_overall"], + mae_entry["mae_per_dim"], + ) + ) + print( + "Inference MAE %s: overall=%g per_dim=%s" + % ( + mae_entry["quantity"], + mae_entry["mae_overall"], + mae_entry["mae_per_dim"], + ), + flush=True, + ) + + for diag_entry in diagnostics_metrics: + print( + "Inference diagnostics %s: bias=%s p90_abs=%s p99_abs=%s high_true_bias=%s" + % ( + diag_entry["quantity"], + diag_entry["bias_per_dim"], + diag_entry["abs_error_p90_per_dim"], + diag_entry["abs_error_p99_per_dim"], + diag_entry["high_true_bias_per_dim"], + ), + flush=True, + ) + + for ihead in range(num_tasks): + name = output_names[ihead] if output_names else f"head{ihead}" + dim = output_dims[ihead] if output_dims else None + _plot_parity_per_dim( + true_values[ihead], + predicted_values[ihead], + name, + dim, + out_dir, + prefix="test", + ) + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/examples/opf/job-andes-preonly.sh b/examples/opf/job-andes-preonly.sh new file mode 100644 index 000000000..86ddabc4d --- /dev/null +++ b/examples/opf/job-andes-preonly.sh @@ -0,0 +1,41 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J HydraGNN +#SBATCH -o job-hydragnn-grid-%j.out +#SBATCH -e job-hydragnn-grid-%j.out +#SBATCH -t 10:00:00 +#SBATCH -p batch +#SBATCH -N 32 +##SBATCH -S 1 + +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +# Load conda environemnt +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-andes.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Andes/hydragnn_venv + +#export python path to use ADIOS2 v.2.10.2 +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Andes/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +#export python path to HydragNN +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +which python +python -c "import numpy; print(numpy.__version__)" +# Core runtime controls +export OMP_NUM_THREADS=7 +export HYDRAGNN_NUM_WORKERS=0 +export HYDRAGNN_USE_VARIABLE_GRAPH_SIZE=1 +export HYDRAGNN_AGGR_BACKEND=mpi +export HYDRAGNN_VALTEST=1 + +cd $HYDRAGNN_ROOT/examples/opf + +srun -N$SLURM_JOB_NUM_NODES -n128 -c7 python -u train_opf_solution_heterogeneous.py --num_groups all --preonly --hdf5 --case_name \ + pglib_opf_case14_ieee pglib_opf_case30_ieee pglib_opf_case57_ieee pglib_opf_case118_ieee pglib_opf_case500_goc pglib_opf_case2000_goc pglib_opf_case4661_sdet pglib_opf_case6470_rte pglib_opf_case10000_goc pglib_opf_case13659_pegase diff --git a/examples/opf/job-frontier-deephyper-hpo-1024.sh b/examples/opf/job-frontier-deephyper-hpo-1024.sh new file mode 100755 index 000000000..71dc0bd07 --- /dev/null +++ b/examples/opf/job-frontier-deephyper-hpo-1024.sh @@ -0,0 +1,105 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J OPF-DeepHyper-1024 +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-hpo-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-hpo-%j.out +#SBATCH -t 12:00:00 +#SBATCH -p batch +#SBATCH -N 1024 + +function cmd() { + echo "$@" + time $@ +} + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +# Load conda environment +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +#export python path to HydraGNN +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +#export python path to use ADIOS2 v.2.10.2 +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +which python +python -c "import adios2; print(adios2.__version__, adios2.__file__)" +python -c "import torch; print(torch.__version__, torch.__file__)" + +module unload darshan-runtime +module list + +echo $LD_LIBRARY_PATH | tr ':' '\n' + +## aws-ofi-rccl plugin settings +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib + +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +export OMP_NUM_THREADS=7 +export HYDRAGNN_AGGR_BACKEND=mpi +export HYDRAGNN_VALTEST=1 +export MPICH_ENV_DISPLAY=0 +export MPICH_VERSION_DISPLAY=0 +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH=/tmp +export PYTHONNOUSERSITE=1 + +## Checking +env | grep ROCM +env | grep ^MI +env | grep ^MPICH +env | grep ^HYDRA + +## OPF HPO settings +export NUM_EPOCH=10 +export BATCH_SIZE=32 + +# DeepHyper configuration +# Each trial uses NNODES_PER_TRIAL nodes; total nodes are split across concurrent trials +export NNODES=$SLURM_JOB_NUM_NODES +export NNODES_PER_TRIAL=16 +export NUM_CONCURRENT_TRIALS=$(( $NNODES / $NNODES_PER_TRIAL )) +export NTOTGPUS=$(( $NNODES * 8 )) +export NGPUS_PER_TRIAL=$(( 8 * $NNODES_PER_TRIAL )) +export NTOT_DEEPHYPER_RANKS=$(( $NTOTGPUS / $NGPUS_PER_TRIAL )) +[ $NTOTGPUS -ne $(($NGPUS_PER_TRIAL*$NUM_CONCURRENT_TRIALS)) ] && echo "ERROR: GPU count mismatch!" + +# DeepHyper variables +export DEEPHYPER_LOG_DIR="deephyper-opf-hpo"-$SLURM_JOB_ID +mkdir -p $DEEPHYPER_LOG_DIR + +cd $HYDRAGNN_ROOT/examples/opf + +which python +python -c "import numpy; print(numpy.__version__)" + +# Use MPNN_TYPE env var if set, otherwise default to all heterogeneous models +[ -z $MPNN_TYPE ] && MPNN_TYPE=HeteroPNA,HeteroSAGE,HeteroGAT,HeteroRGAT,HeteroHGT,HeteroHEAT + +cmd python -u $HYDRAGNN_ROOT/examples/opf/opf_deephyper_hpo.py \ + --mpnn_type=$MPNN_TYPE \ + --max_evals=200 diff --git a/examples/opf/job-frontier-deephyper-hpo-pna.sh b/examples/opf/job-frontier-deephyper-hpo-pna.sh new file mode 100755 index 000000000..b057a3cd7 --- /dev/null +++ b/examples/opf/job-frontier-deephyper-hpo-pna.sh @@ -0,0 +1,107 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J OPF-HPO-PNA +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-hpo-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-hpo-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 128 + +function cmd() { + echo "$@" + time $@ +} + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +# Load conda environment +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +#export python path to HydraGNN +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +#export python path to use ADIOS2 v.2.10.2 +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +which python +python -c "import adios2; print(adios2.__version__, adios2.__file__)" +python -c "import torch; print(torch.__version__, torch.__file__)" + +module unload darshan-runtime +module list + +echo $LD_LIBRARY_PATH | tr ':' '\n' + +## aws-ofi-rccl plugin settings +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib + +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +export OMP_NUM_THREADS=7 +export HYDRAGNN_AGGR_BACKEND=mpi +export HYDRAGNN_VALTEST=1 +export MPICH_ENV_DISPLAY=0 +export MPICH_VERSION_DISPLAY=0 +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH=/tmp +export PYTHONNOUSERSITE=1 + +## Checking +env | grep ROCM +env | grep ^MI +env | grep ^MPICH +env | grep ^HYDRA + +## OPF HPO settings +export NUM_EPOCH=10 +export BATCH_SIZE=32 + +# DeepHyper configuration +export NNODES=$SLURM_JOB_NUM_NODES +export NNODES_PER_TRIAL=16 +export NUM_CONCURRENT_TRIALS=$(( $NNODES / $NNODES_PER_TRIAL )) +export NTOTGPUS=$(( $NNODES * 8 )) +export NGPUS_PER_TRIAL=$(( 8 * $NNODES_PER_TRIAL )) +export NTOT_DEEPHYPER_RANKS=$(( $NTOTGPUS / $NGPUS_PER_TRIAL )) +[ $NTOTGPUS -ne $(($NGPUS_PER_TRIAL*$NUM_CONCURRENT_TRIALS)) ] && echo "ERROR: GPU count mismatch!" + +export DEEPHYPER_LOG_DIR="deephyper-opf-hpo"-$SLURM_JOB_ID +mkdir -p $DEEPHYPER_LOG_DIR + +cd $HYDRAGNN_ROOT/examples/opf + +which python +python -c "import numpy; print(numpy.__version__)" + +# HeteroPNA-only HPO with constrained search space around the known-working config +# Working config: hidden_dim=64, num_conv_layers=4, lr=0.001 +# Constrain to safe memory range: hidden_dim 32-64, layers 2-4 +cmd python -u $HYDRAGNN_ROOT/examples/opf/opf_deephyper_hpo.py \ + --mpnn_type=HeteroPNA \ + --hidden_dim_range=32,64 \ + --num_conv_layers_range=2,4 \ + --learning_rate_range=1e-4,1e-2 \ + --max_evals=50 diff --git a/examples/opf/job-frontier-deephyper-hpo.sh b/examples/opf/job-frontier-deephyper-hpo.sh new file mode 100644 index 000000000..573d6c168 --- /dev/null +++ b/examples/opf/job-frontier-deephyper-hpo.sh @@ -0,0 +1,104 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J OPF-DeepHyper +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-hpo-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-hpo-%j.out +#SBATCH -t 06:00:00 +#SBATCH -p batch +#SBATCH -N 128 + +function cmd() { + echo "$@" + time $@ +} + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +# Load conda environment +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +#export python path to HydraGNN +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +#export python path to use ADIOS2 v.2.10.2 +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +which python +python -c "import adios2; print(adios2.__version__, adios2.__file__)" +python -c "import torch; print(torch.__version__, torch.__file__)" + +module unload darshan-runtime +module list + +echo $LD_LIBRARY_PATH | tr ':' '\n' + +## aws-ofi-rccl plugin settings +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib + +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +export OMP_NUM_THREADS=7 +export HYDRAGNN_AGGR_BACKEND=mpi +export HYDRAGNN_VALTEST=1 +export MPICH_ENV_DISPLAY=0 +export MPICH_VERSION_DISPLAY=0 +export MIOPEN_DISABLE_CACHE=1 +export MIOPEN_USER_DB_PATH=/tmp +export PYTHONNOUSERSITE=1 + +## Checking +env | grep ROCM +env | grep ^MI +env | grep ^MPICH +env | grep ^HYDRA + +## OPF HPO settings +export NUM_EPOCH=10 +export BATCH_SIZE=32 + +# DeepHyper configuration +# Each trial uses NNODES_PER_TRIAL nodes; total nodes are split across concurrent trials +export NNODES=$SLURM_JOB_NUM_NODES +export NNODES_PER_TRIAL=16 +export NUM_CONCURRENT_TRIALS=$(( $NNODES / $NNODES_PER_TRIAL )) +export NTOTGPUS=$(( $NNODES * 8 )) +export NGPUS_PER_TRIAL=$(( 8 * $NNODES_PER_TRIAL )) +export NTOT_DEEPHYPER_RANKS=$(( $NTOTGPUS / $NGPUS_PER_TRIAL )) +[ $NTOTGPUS -ne $(($NGPUS_PER_TRIAL*$NUM_CONCURRENT_TRIALS)) ] && echo "ERROR: GPU count mismatch!" + +# DeepHyper variables +export DEEPHYPER_LOG_DIR="deephyper-opf-hpo"-$SLURM_JOB_ID +mkdir -p $DEEPHYPER_LOG_DIR + +cd $HYDRAGNN_ROOT/examples/opf + +which python +python -c "import numpy; print(numpy.__version__)" + +# Use MPNN_TYPE env var if set, otherwise default to all heterogeneous models +[ -z $MPNN_TYPE ] && MPNN_TYPE=HeteroPNA,HeteroSAGE,HeteroGAT,HeteroRGAT,HeteroHGT,HeteroHEAT + +cmd python -u $HYDRAGNN_ROOT/examples/opf/opf_deephyper_hpo.py \ + --mpnn_type=$MPNN_TYPE diff --git a/examples/opf/job-frontier-inference.sh b/examples/opf/job-frontier-inference.sh new file mode 100644 index 000000000..c41764ec8 --- /dev/null +++ b/examples/opf/job-frontier-inference.sh @@ -0,0 +1,92 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J HydraGNN +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/job-hydragnn-grid-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/job-hydragnn-grid-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 5 #16 +##SBATCH -S 1 + +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +function cmd() { + echo "$@" + time $@ +} + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +# Load conda environemnt +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +#export python path to HydragNN +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +#export python path to use ADIOS2 v.2.10.2 +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +which python +python -c "import adios2; print(adios2.__version__, adios2.__file__)" +python -c "import torch; print(torch.__version__, torch.__file__)" + +module unload darshan-runtime +module list + + +echo $LD_LIBRARY_PATH | tr ':' '\n' + +## aws-ofi-rccl plugin settings +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib + +export FI_MR_CACHE_MONITOR=kdreg2 # Required to avoid a deadlock. +export FI_CXI_DEFAULT_CQ_SIZE=131072 # Ask the network stack to allocate additional space to process message completions. +export FI_CXI_DEFAULT_TX_SIZE=2048 # Ask the network stack to allocate additional space to hold pending outgoing messages. +export FI_CXI_RX_MATCH_MODE=hybrid # Allow the network stack to transition to software mode if necessary. +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" # Typically improves performance, but remove this setting if you encounter a hang/crash. +export NCCL_CROSS_NIC=1 # On large systems, this NCCL setting has been found to improve performance +export NCCL_SOCKET_IFNAME=hsn0 # NCCL/RCCL will use the high speed network to coordinate startup. +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 # Use high priority stream for the NCCL/RCCL Communicator. +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +# below are optional to debug RCCL stuff +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=INIT + +# The following have been found to help avoid hangs, but are not yet +# documented elsewhere +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + + +## Checking +env | grep ROCM +env | grep ^MI +env | grep ^MPICH +env | grep ^HYDRA + +cd $HYDRAGNN_ROOT/examples/opf + +which python +python -c "import numpy; print(numpy.__version__)" + +#srun -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 --gpus-per-task=1 --gpu-bind=closest python -u train_mlip.py --preonly --hdf5 --ddstore +#srun -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 --gpus-per-task=1 --gpu-bind=closest python -u train_opf_heterogeneous.py --preonly --hdf5 +#srun -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 --gpus-per-task=1 --gpu-bind=closest python -u train_opf_solution_heterogeneous.py --preonly --hdf5 --case_name pglib_opf_case118_ieee pglib_opf_case14_ieee pglib_opf_case2000_goc pglib_opf_case30_ieee pglib_opf_case500_goc pglib_opf_case57_ieee +srun -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 --gpus-per-task=1 --gpu-bind=closest python -u infer_opf_solution_heterogeneous.py --hdf5 --modelname OPF_Solution_Hetero diff --git a/examples/opf/job-frontier-preonly.sh b/examples/opf/job-frontier-preonly.sh new file mode 100644 index 000000000..da1ca8f99 --- /dev/null +++ b/examples/opf/job-frontier-preonly.sh @@ -0,0 +1,83 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J HydraGNN +#SBATCH -o job-hydragnn-grid-%j.out +#SBATCH -e job-hydragnn-grid-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 16 +##SBATCH -C nvme +##SBATCH -S 1 + +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +# Load conda environemnt +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +#export python path to HydragNN +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +#export python path to use ADIOS2 v.2.10.2 +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +which python +python -c "import adios2; print(adios2.__version__, adios2.__file__)" +python -c "import torch; print(torch.__version__, torch.__file__)" + +module unload darshan-runtime +module list + + +## aws-ofi-rccl plugin settings +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib + +export FI_MR_CACHE_MONITOR=kdreg2 # Required to avoid a deadlock. +export FI_CXI_DEFAULT_CQ_SIZE=131072 # Ask the network stack to allocate additional space to process message completions. +export FI_CXI_DEFAULT_TX_SIZE=2048 # Ask the network stack to allocate additional space to hold pending outgoing messages. +export FI_CXI_RX_MATCH_MODE=hybrid # Allow the network stack to transition to software mode if necessary. +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" # Typically improves performance, but remove this setting if you encounter a hang/crash. +export NCCL_CROSS_NIC=1 # On large systems, this NCCL setting has been found to improve performance +export NCCL_SOCKET_IFNAME=hsn0 # NCCL/RCCL will use the high speed network to coordinate startup. +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 # Use high priority stream for the NCCL/RCCL Communicator. +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +# below are optional to debug RCCL stuff +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=INIT + +# The following have been found to help avoid hangs, but are not yet +# documented elsewhere +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + + +# Diagnostics +export HYDRAGNN_DIAG=1 +export HYDRAGNN_DIAG_RANK=1 + +cd $HYDRAGNN_ROOT/examples/opf + +srun --export=ALL,HYDRAGNN_DIAG=1,HYDRAGNN_DIAG_RANK=1 \ + -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_solution_heterogeneous.py --num_groups all --preonly --hdf5 --case_name \ + pglib_opf_case14_ieee pglib_opf_case30_ieee pglib_opf_case57_ieee pglib_opf_case118_ieee pglib_opf_case500_goc pglib_opf_case2000_goc pglib_opf_case4661_sdet pglib_opf_case6470_rte pglib_opf_case10000_goc pglib_opf_case13659_pegase + diff --git a/examples/opf/job-frontier-train-heat-compare.sh b/examples/opf/job-frontier-train-heat-compare.sh new file mode 100644 index 000000000..dbbdd9f39 --- /dev/null +++ b/examples/opf/job-frontier-train-heat-compare.sh @@ -0,0 +1,131 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J OPF-HEAT-AB +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-heat-ab-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-heat-ab-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 16 + +set -euo pipefail + +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN +INPUTFILE=${INPUTFILE:-logs/opf_hpo_4249563_0.28/config.json} +MODELNAME=${MODELNAME:-OPF_Solution_Hetero} +BASELINE_LOG_NAME=${BASELINE_LOG_NAME:-opf_heat_hpo_best_baseline} +PHYSICS_LOG_NAME=${PHYSICS_LOG_NAME:-opf_heat_hpo_best_physics} + +# Domain-loss weights for the physics-informed run. +# Use two separate index variables to avoid fragile string-splitting. +DOMAIN_SMOOTHNESS_WEIGHT=${DOMAIN_SMOOTHNESS_WEIGHT:-0.001} +DOMAIN_TRANSFORMER_SMOOTHNESS_WEIGHT=${DOMAIN_TRANSFORMER_SMOOTHNESS_WEIGHT:-0.001} +DOMAIN_VOLTAGE_BOUND_WEIGHT=${DOMAIN_VOLTAGE_BOUND_WEIGHT:-0.01} +VMIN_IDX=${VMIN_IDX:-2} +VMAX_IDX=${VMAX_IDX:-3} +# voltage_output_index=1: Vm is at index 1 in bus targets [Va(0), Vm(1)] +DOMAIN_VOLTAGE_OUTPUT_INDEX=${DOMAIN_VOLTAGE_OUTPUT_INDEX:-1} +# va_output_index=0: Va is at index 0 in bus targets +DOMAIN_VA_OUTPUT_INDEX=${DOMAIN_VA_OUTPUT_INDEX:-0} +# angle difference limit and DC thermal limit penalties (new deeper physics terms) +DOMAIN_ANGLE_DIFF_WEIGHT=${DOMAIN_ANGLE_DIFF_WEIGHT:-0.001} +DOMAIN_LINE_FLOW_WEIGHT=${DOMAIN_LINE_FLOW_WEIGHT:-0.001} + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL=PHB +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +cd $HYDRAGNN_ROOT/examples/opf + +which python +python -c "import torch; print(torch.__version__, torch.__file__)" +python -c "import numpy; print(numpy.__version__)" + +run_train() { + local log_name=$1 + shift + echo "" + echo "=====================================================================" + echo " STARTING: INPUTFILE=$INPUTFILE MODELNAME=$MODELNAME LOG=$log_name" + echo " $(date)" + echo "=====================================================================" + srun --export=ALL,HYDRAGNN_DIAG=1,HYDRAGNN_DIAG_RANK=0 \ + -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_solution_heterogeneous.py \ + --hdf5 \ + --inputfile "$INPUTFILE" \ + --modelname "$MODELNAME" \ + --log "$log_name" \ + "$@" + echo "" + echo "=====================================================================" + echo " COMPLETED: $log_name $(date)" + echo "=====================================================================" + # Sanity-check: verify the run produced a log file. + local runlog="logs/$log_name/run.log" + if [[ -f "$runlog" ]]; then + echo " Log found: $runlog" + else + echo " WARNING: expected log not found at $runlog" + fi +} + +# ── Run 1: Baseline (standard loss) ────────────────────────────────────────── +run_train "$BASELINE_LOG_NAME" --disable_domain_loss + +# ── Run 2: Physics-informed (domain loss enabled) ──────────────────────────── +run_train "$PHYSICS_LOG_NAME" \ + --enable_domain_loss \ + --domain_loss_smoothness_weight "$DOMAIN_SMOOTHNESS_WEIGHT" \ + --domain_loss_transformer_smoothness_weight "$DOMAIN_TRANSFORMER_SMOOTHNESS_WEIGHT" \ + --domain_loss_voltage_bound_weight "$DOMAIN_VOLTAGE_BOUND_WEIGHT" \ + --domain_loss_voltage_bound_feature_indices "$VMIN_IDX" "$VMAX_IDX" \ + --domain_loss_voltage_output_index "$DOMAIN_VOLTAGE_OUTPUT_INDEX" \ + --domain_loss_va_output_index "$DOMAIN_VA_OUTPUT_INDEX" \ + --domain_loss_angle_diff_weight "$DOMAIN_ANGLE_DIFF_WEIGHT" \ + --domain_loss_line_flow_weight "$DOMAIN_LINE_FLOW_WEIGHT" + +# ── Final summary ───────────────────────────────────────────────────────────── +echo "" +echo "=====================================================================" +echo " BOTH RUNS COMPLETE $(date)" +echo " Baseline log : logs/$BASELINE_LOG_NAME/run.log" +echo " Physics log : logs/$PHYSICS_LOG_NAME/run.log" +echo "" +echo " To compare results run:" +echo " python3 compare_heat_runs.py \\" +echo " logs/$BASELINE_LOG_NAME/run.log \\" +echo " logs/$PHYSICS_LOG_NAME/run.log" +echo " For JSON output add --json" +echo "=====================================================================" diff --git a/examples/opf/job-frontier-train-heat-physics.sh b/examples/opf/job-frontier-train-heat-physics.sh new file mode 100644 index 000000000..1b878fc16 --- /dev/null +++ b/examples/opf/job-frontier-train-heat-physics.sh @@ -0,0 +1,120 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J OPF-HEAT-PI +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-heat-pi-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/job-opf-heat-pi-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 16 + +set -euo pipefail + +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN +INPUTFILE=${INPUTFILE:-logs/opf_hpo_4249563_0.28/config.json} +MODELNAME=${MODELNAME:-OPF_Solution_Hetero} +PHYSICS_LOG_NAME=${PHYSICS_LOG_NAME:-opf_heat_hpo_best_physics} + +# Domain-loss weights. +# Weights are set to ~10% of the expected converged task loss (~0.009) divided +# by the number of active penalty terms, so the total physics contribution +# stays well below the data-driven MSE and acts as a soft regulariser rather +# than a competing objective. +DOMAIN_VOLTAGE_BOUND_WEIGHT=${DOMAIN_VOLTAGE_BOUND_WEIGHT:-0.001} +VMIN_IDX=${VMIN_IDX:-2} +VMAX_IDX=${VMAX_IDX:-3} +# voltage_output_index=1: Vm is at bus_pred[:,1] (Va=0, Vm=1 in OPFDataset schema) +DOMAIN_VOLTAGE_OUTPUT_INDEX=${DOMAIN_VOLTAGE_OUTPUT_INDEX:-1} +# va_output_index=0: Va is at bus_pred[:,0] +DOMAIN_VA_OUTPUT_INDEX=${DOMAIN_VA_OUTPUT_INDEX:-0} +DOMAIN_ANGLE_DIFF_WEIGHT=${DOMAIN_ANGLE_DIFF_WEIGHT:-0.0001} +DOMAIN_LINE_FLOW_WEIGHT=${DOMAIN_LINE_FLOW_WEIGHT:-0.0001} +# Curriculum scheduling: warmup then linear ramp before full domain-loss weight. +# With num_epoch=10: epochs 0-2 task-loss only, 3-5 linear ramp, 6-9 full weight. +DOMAIN_WARMUP_EPOCHS=${DOMAIN_WARMUP_EPOCHS:-3} +DOMAIN_RAMP_EPOCHS=${DOMAIN_RAMP_EPOCHS:-3} + +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +export PYTHONPATH=$HYDRAGNN_ROOT:${PYTHONPATH:-} +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +module unload darshan-runtime + +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib +export FI_MR_CACHE_MONITOR=kdreg2 +export FI_CXI_DEFAULT_CQ_SIZE=131072 +export FI_CXI_DEFAULT_TX_SIZE=2048 +export FI_CXI_RX_MATCH_MODE=hybrid +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL=PHB +export NCCL_CROSS_NIC=1 +export NCCL_SOCKET_IFNAME=hsn0 +export NCCL_NET="AWS Libfabric" +export TORCH_NCCL_HIGH_PRIORITY=1 +export GPU_MAX_HW_QUEUES=2 +export HSA_FORCE_FINE_GRAIN_PCIE=1 +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +cd $HYDRAGNN_ROOT/examples/opf + +which python3 +python3 -c "import torch; print(torch.__version__, torch.__file__)" +python3 -c "import numpy; print(numpy.__version__)" + +echo "" +echo "=====================================================================" +echo " STARTING physics-informed run" +echo " INPUTFILE=$INPUTFILE MODELNAME=$MODELNAME LOG=$PHYSICS_LOG_NAME" +echo " $(date)" +echo " Penalties: voltage_bound=$DOMAIN_VOLTAGE_BOUND_WEIGHT (Vm@idx $DOMAIN_VOLTAGE_OUTPUT_INDEX, [vmin@$VMIN_IDX, vmax@$VMAX_IDX])" +echo " angle_diff=$DOMAIN_ANGLE_DIFF_WEIGHT (Va@idx $DOMAIN_VA_OUTPUT_INDEX)" +echo " line_flow=$DOMAIN_LINE_FLOW_WEIGHT" echo " Curriculum: warmup=$DOMAIN_WARMUP_EPOCHS epochs, ramp=$DOMAIN_RAMP_EPOCHS epochs"echo "=====================================================================" + +srun --export=ALL,HYDRAGNN_DIAG=1,HYDRAGNN_DIAG_RANK=0 \ + -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python3 -u train_opf_solution_heterogeneous.py \ + --hdf5 \ + --inputfile "$INPUTFILE" \ + --modelname "$MODELNAME" \ + --log "$PHYSICS_LOG_NAME" \ + --enable_domain_loss \ + --domain_loss_voltage_bound_weight "$DOMAIN_VOLTAGE_BOUND_WEIGHT" \ + --domain_loss_voltage_bound_feature_indices "$VMIN_IDX" "$VMAX_IDX" \ + --domain_loss_voltage_output_index "$DOMAIN_VOLTAGE_OUTPUT_INDEX" \ + --domain_loss_va_output_index "$DOMAIN_VA_OUTPUT_INDEX" \ + --domain_loss_angle_diff_weight "$DOMAIN_ANGLE_DIFF_WEIGHT" \ + --domain_loss_line_flow_weight "$DOMAIN_LINE_FLOW_WEIGHT" \ + --domain_loss_warmup_epochs "$DOMAIN_WARMUP_EPOCHS" \ + --domain_loss_ramp_epochs "$DOMAIN_RAMP_EPOCHS" + +echo "" +echo "=====================================================================" +echo " COMPLETED: $PHYSICS_LOG_NAME $(date)" +echo "=====================================================================" + +runlog="logs/$PHYSICS_LOG_NAME/run.log" +if [[ -f "$runlog" ]]; then + echo " Log found: $runlog" +else + echo " WARNING: expected log not found at $runlog" +fi + +echo "" +echo " To compare with baseline:" +echo " python3 compare_heat_runs.py \\" +echo " logs/opf_heat_hpo_best_baseline/run.log \\" +echo " logs/$PHYSICS_LOG_NAME/run.log" diff --git a/examples/opf/job-frontier-train-heat.sh b/examples/opf/job-frontier-train-heat.sh new file mode 100755 index 000000000..5f9ee0d4a --- /dev/null +++ b/examples/opf/job-frontier-train-heat.sh @@ -0,0 +1,92 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J HydraGNN +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/job-hydragnn-grid-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/job-hydragnn-grid-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 16 +##SBATCH -C nvme +##SBATCH -S 1 + +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +function cmd() { + echo "$@" + time $@ +} + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +# Load conda environemnt +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +#export python path to HydragNN +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +#export python path to use ADIOS2 v.2.10.2 +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +which python +python -c "import adios2; print(adios2.__version__, adios2.__file__)" +python -c "import torch; print(torch.__version__, torch.__file__)" + +module unload darshan-runtime +module list + + +echo $LD_LIBRARY_PATH | tr ':' '\n' + +## aws-ofi-rccl plugin settings +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib + +export FI_MR_CACHE_MONITOR=kdreg2 # Required to avoid a deadlock. +export FI_CXI_DEFAULT_CQ_SIZE=131072 # Ask the network stack to allocate additional space to process message completions. +export FI_CXI_DEFAULT_TX_SIZE=2048 # Ask the network stack to allocate additional space to hold pending outgoing messages. +export FI_CXI_RX_MATCH_MODE=hybrid # Allow the network stack to transition to software mode if necessary. +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" # Typically improves performance, but remove this setting if you encounter a hang/crash. +export NCCL_CROSS_NIC=1 # On large systems, this NCCL setting has been found to improve performance +export NCCL_SOCKET_IFNAME=hsn0 # NCCL/RCCL will use the high speed network to coordinate startup. +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 # Use high priority stream for the NCCL/RCCL Communicator. +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +# below are optional to debug RCCL stuff +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=INIT + +# The following have been found to help avoid hangs, but are not yet +# documented elsewhere +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +## Checking +env | grep ROCM +env | grep ^MI +env | grep ^MPICH +env | grep ^HYDRA + +cd $HYDRAGNN_ROOT/examples/opf + +which python +python -c "import numpy; print(numpy.__version__)" + +srun --export=ALL,HYDRAGNN_DIAG=1,HYDRAGNN_DIAG_RANK=0 \ + -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_solution_heterogeneous.py --hdf5 --inputfile opf_solution_heterogeneous_heat.json --modelname OPF_Solution_Hetero diff --git a/examples/opf/job-frontier-train.sh b/examples/opf/job-frontier-train.sh new file mode 100644 index 000000000..46484d3a1 --- /dev/null +++ b/examples/opf/job-frontier-train.sh @@ -0,0 +1,92 @@ +#!/bin/bash +#SBATCH -A LRN078 +#SBATCH -J HydraGNN +#SBATCH -o /lustre/orion/lrn078/proj-shared/HydraGNN/job-hydragnn-grid-%j.out +#SBATCH -e /lustre/orion/lrn078/proj-shared/HydraGNN/job-hydragnn-grid-%j.out +#SBATCH -t 02:00:00 +#SBATCH -p batch +#SBATCH -q debug +#SBATCH -N 16 +##SBATCH -C nvme +##SBATCH -S 1 + +export all_proxy=socks://proxy.ccs.ornl.gov:3128/ +export ftp_proxy=ftp://proxy.ccs.ornl.gov:3128/ +export http_proxy=http://proxy.ccs.ornl.gov:3128/ +export https_proxy=http://proxy.ccs.ornl.gov:3128/ +export no_proxy='localhost,127.0.0.0/8,*.ccs.ornl.gov' + +function cmd() { + echo "$@" + time $@ +} + +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN + +# Load conda environemnt +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv + +#export python path to HydragNN +export PYTHONPATH=$HYDRAGNN_ROOT:$PYTHONPATH + +#export python path to use ADIOS2 v.2.10.2 +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH + +which python +python -c "import adios2; print(adios2.__version__, adios2.__file__)" +python -c "import torch; print(torch.__version__, torch.__file__)" + +module unload darshan-runtime +module list + + +echo $LD_LIBRARY_PATH | tr ':' '\n' + +## aws-ofi-rccl plugin settings +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib + +export FI_MR_CACHE_MONITOR=kdreg2 # Required to avoid a deadlock. +export FI_CXI_DEFAULT_CQ_SIZE=131072 # Ask the network stack to allocate additional space to process message completions. +export FI_CXI_DEFAULT_TX_SIZE=2048 # Ask the network stack to allocate additional space to hold pending outgoing messages. +export FI_CXI_RX_MATCH_MODE=hybrid # Allow the network stack to transition to software mode if necessary. +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" # Typically improves performance, but remove this setting if you encounter a hang/crash. +export NCCL_CROSS_NIC=1 # On large systems, this NCCL setting has been found to improve performance +export NCCL_SOCKET_IFNAME=hsn0 # NCCL/RCCL will use the high speed network to coordinate startup. +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 # Use high priority stream for the NCCL/RCCL Communicator. +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +# below are optional to debug RCCL stuff +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=INIT + +# The following have been found to help avoid hangs, but are not yet +# documented elsewhere +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 + +## Checking +env | grep ROCM +env | grep ^MI +env | grep ^MPICH +env | grep ^HYDRA + +cd $HYDRAGNN_ROOT/examples/opf + +which python +python -c "import numpy; print(numpy.__version__)" + +srun --export=ALL,HYDRAGNN_DIAG=1,HYDRAGNN_DIAG_RANK=0 \ + -N$SLURM_JOB_NUM_NODES -n$((SLURM_JOB_NUM_NODES*8)) -c7 \ + --gpus-per-task=1 --gpu-bind=closest \ + python -u train_opf_solution_heterogeneous.py --hdf5 diff --git a/examples/opf/opf_deephyper_hpo.py b/examples/opf/opf_deephyper_hpo.py new file mode 100644 index 000000000..ba0d3d377 --- /dev/null +++ b/examples/opf/opf_deephyper_hpo.py @@ -0,0 +1,299 @@ +import os, sys + +import math +import torch + +torch.backends.cudnn.enabled = False + +# FIX random seed +random_state = 0 +torch.manual_seed(random_state) + +import pandas as pd +import subprocess +import re +import argparse +import glob + +pd.options.display.max_columns = None +pd.options.display.max_rows = None +pd.options.display.width = None + +# Retrieve constants +NNODES = int(os.environ["NNODES"]) +NTOTGPUS = int(os.environ["NTOTGPUS"]) +NNODES_PER_TRIAL = int(os.environ["NNODES_PER_TRIAL"]) +NGPUS_PER_TRIAL = int(os.environ["NGPUS_PER_TRIAL"]) +NUM_CONCURRENT_TRIALS = int(os.environ["NUM_CONCURRENT_TRIALS"]) +NTOT_DEEPHYPER_RANKS = int(os.environ["NTOT_DEEPHYPER_RANKS"]) +OMP_NUM_THREADS = int(os.environ["OMP_NUM_THREADS"]) +DEEPHYPER_LOG_DIR = os.environ["DEEPHYPER_LOG_DIR"] +SLURM_JOB_ID = os.environ["SLURM_JOB_ID"] + +# OPF-specific environment variables +NUM_EPOCH = int(os.environ["NUM_EPOCH"]) +BATCH_SIZE = int(os.environ["BATCH_SIZE"]) + + +def to_float(x): + x = x.lower() + if x == "nan": + return math.nan + if x in ("inf", "+inf"): + return math.inf + if x == "-inf": + return -math.inf + return float(x) + + +def run(trial, dequed=None): + os.makedirs(f"{DEEPHYPER_LOG_DIR}", exist_ok=True) + f = open(f"{DEEPHYPER_LOG_DIR}/output-{trial.id}.txt", "w") + python_exe = sys.executable + python_script = os.path.join( + os.path.dirname(__file__), "train_opf_solution_heterogeneous.py" + ) + + log_name = f"opf_hpo_{SLURM_JOB_ID}_{trial.id}" + master_addr = f"HYDRAGNN_MASTER_ADDR={dequed[0]}" + nodelist = ",".join(dequed) + + prefix = " ".join( + [ + f"srun", + f"-N {NNODES_PER_TRIAL} -n {NGPUS_PER_TRIAL}", + f"--ntasks-per-node=8 --gpus-per-node=8", + f"--cpus-per-task {OMP_NUM_THREADS} --threads-per-core 1 --cpu-bind threads", + f"--gpus-per-task=1 --gpu-bind=closest", + f"--export=ALL,{master_addr}", + f"--nodelist={nodelist}", + f"--output {DEEPHYPER_LOG_DIR}/output_{SLURM_JOB_ID}_{trial.id}.txt", + f"--error {DEEPHYPER_LOG_DIR}/error_{SLURM_JOB_ID}_{trial.id}.txt", + ] + ) + + command = " ".join( + [ + prefix, + python_exe, + "-u", + python_script, + f"--inputfile=opf_solution_heterogeneous.json", + f"--hdf5", + f"--num_epoch={NUM_EPOCH}", + f"--batch_size={BATCH_SIZE}", + f"--log={log_name}", + f"--mpnn_type={trial.parameters['mpnn_type']}", + f"--hidden_dim={trial.parameters['hidden_dim']}", + f"--num_conv_layers={trial.parameters['num_conv_layers']}", + f"--learning_rate={trial.parameters['learning_rate']}", + ] + ) + + print("Command = ", command, flush=True, file=f) + + output = -math.inf + num_pattern = r"[-+]?(?:\d*\.\d+|\d+\.?)(?:[eE][-+]?\d+)?|[-+]?(?:inf|nan)" + try: + result = subprocess.check_output(command, shell=True, stderr=subprocess.STDOUT) + fout = open(f"{DEEPHYPER_LOG_DIR}/error_{SLURM_JOB_ID}_{trial.id}.txt", "r") + while True: + line = fout.readline() + if "Tasks Val Loss:" in line: + nums = re.findall(num_pattern, line, flags=re.IGNORECASE) + # nums[0] is the rank prefix (e.g. "0:"), nums[1] is the actual loss + if len(nums) >= 2: + val = -to_float(nums[1]) + print( + f"Val loss: {-val}", + flush=True, + file=f, + ) + # Keep the best (minimum) val loss across epochs + if val > output: + output = val + if not line: + break + fout.close() + except subprocess.CalledProcessError as cpe: + # If the trial was killed by walltime, epochs may have completed. + # Try to extract the best val loss from whatever was logged. + print(f"Trial failed with exit code {cpe.returncode}", flush=True, file=f) + error_file = f"{DEEPHYPER_LOG_DIR}/error_{SLURM_JOB_ID}_{trial.id}.txt" + if os.path.exists(error_file): + fout = open(error_file, "r") + for line in fout: + if "Tasks Val Loss:" in line: + nums = re.findall(num_pattern, line, flags=re.IGNORECASE) + if len(nums) >= 2: + val = -to_float(nums[1]) + print( + f"Val loss (from partial run): {-val}", + flush=True, + file=f, + ) + if val > output: + output = val + fout.close() + except Exception as excp: + print(excp, flush=True, file=f) + output = -math.inf + + print(f"Best val loss (min across epochs): {-output}", flush=True, file=f) + print("Output:", output, flush=True, file=f) + objective = output + print(objective, flush=True, file=f) + metadata = {"mpnn_type": trial.parameters["mpnn_type"]} + f.close() + + return {"objective": objective, "metadata": metadata} + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "--mpnn_type", + type=str, + default="HeteroPNA,HeteroSAGE,HeteroGAT,HeteroRGAT,HeteroHGT,HeteroHEAT", + ) + parser.add_argument( + "--max_evals", + type=int, + default=100, + help="Number of max evaluations for HPO search", + ) + parser.add_argument( + "--hidden_dim_range", + type=str, + default="32,256", + help="min,max for hidden_dim (e.g. '32,64')", + ) + parser.add_argument( + "--num_conv_layers_range", + type=str, + default="2,6", + help="min,max for num_conv_layers (e.g. '2,4')", + ) + parser.add_argument( + "--learning_rate_range", + type=str, + default="1e-5,1e-2", + help="min,max for learning_rate (e.g. '1e-4,1e-2')", + ) + parser.add_argument( + "--preload_dirs", + type=str, + default="", + help="Comma-separated list of DeepHyper log dirs to warm-start from (e.g. 'opf_hpo_HeteroSAGE-4466182')", + ) + args = parser.parse_args() + mpnn_type_list = args.mpnn_type.split(",") + + log_name = f"opf_hpo-{SLURM_JOB_ID}" + if len(mpnn_type_list) == 1: + log_name = f"opf_hpo_{mpnn_type_list[0]}-{SLURM_JOB_ID}" + + from deephyper.hpo import HpProblem, CBO + from deephyper.evaluator import ProcessPoolEvaluator, queued + from hydragnn.utils.hpo.deephyper import read_node_list + + hd_min, hd_max = (int(x) for x in args.hidden_dim_range.split(",")) + cl_min, cl_max = (int(x) for x in args.num_conv_layers_range.split(",")) + lr_min, lr_max = (float(x) for x in args.learning_rate_range.split(",")) + + hyperparameters = dict() + hyperparameters["mpnn_type"] = mpnn_type_list + hyperparameters["learning_rate"] = (lr_min, lr_max) + hyperparameters["hidden_dim"] = (hd_min, hd_max) + hyperparameters["num_conv_layers"] = (cl_min, cl_max) + + ## Create HPO problem with the defined hyperparameters + problem = HpProblem() + for k, v in hyperparameters.items(): + problem.add_hyperparameter(v, k) + + # Create the node queue + queue, _ = read_node_list() + print("The queue:", queue, len(queue)) + print("NNODES_PER_TRIAL", NNODES_PER_TRIAL) + print("NUM_CONCURRENT_TRIALS", NUM_CONCURRENT_TRIALS) + print("NGPUS_PER_TRIAL", NGPUS_PER_TRIAL) + print("NTOTGPUS", NTOTGPUS) + print(NTOTGPUS, NGPUS_PER_TRIAL, NTOTGPUS // NGPUS_PER_TRIAL, len(queue)) + + evaluator = queued(ProcessPoolEvaluator)( + run, + num_workers=NUM_CONCURRENT_TRIALS, + queue=queue, + queue_pop_per_task=NNODES_PER_TRIAL, + ) + + search = CBO( + problem, + acq_func="UCB", + multi_point_strategy="cl_min", + random_state=42, + log_dir=log_name, + ) + + ## Preload results from previous runs + opf_dir = "opf_hpo" + if len(mpnn_type_list) == 1: + opf_dir = f"opf_hpo_{mpnn_type_list[0]}" + + # Build list of directories to search for prior results + preload_dirs = [opf_dir] + if args.preload_dirs: + preload_dirs += [d.strip() for d in args.preload_dirs.split(",") if d.strip()] + print(f"Warm-start: searching for prior results in: {preload_dirs}") + + df_list = list() + files = [] + for d in preload_dirs: + files += glob.glob(os.path.join(d, "*.csv")) + for fname in files: + try: + df = pd.read_csv(fname, header=0) + total_rows = len(df) + df["objective"] = pd.to_numeric(df["objective"], errors="coerce") + df = df.dropna(subset=["objective"]) + valid_rows = list() + for i in range(len(df)): + try: + search.fit_surrogate(df.iloc[i : i + 1]) + valid_rows.append(i) + except: + continue + print(f"Checking {fname}: total {total_rows}, valid {len(valid_rows)}") + df = df.iloc[valid_rows] + df_list.append(df) + except Exception as excp: + print(f"Error loading {fname}:", excp) + + ## Create a clean search object again + search = CBO( + problem, + acq_func="UCB", + multi_point_strategy="cl_min", + random_state=42, + log_dir=log_name, + ) + + if len(df_list) > 0: + try: + preloaded_results = pd.concat(df_list, ignore_index=True) + print( + f"Loaded {len(preloaded_results)} preloaded results from {len(df_list)} files." + ) + search.fit_surrogate(preloaded_results) + except Exception as excp: + print("Error in loading preloaded results:", excp) + + timeout = None + results = search.search(evaluator, max_evals=args.max_evals, timeout=timeout) + print(results) + + sys.exit(0) diff --git a/examples/opf/opf_graph_output_heterogeneous.json b/examples/opf/opf_graph_output_heterogeneous.json new file mode 100644 index 000000000..1dce62c59 --- /dev/null +++ b/examples/opf/opf_graph_output_heterogeneous.json @@ -0,0 +1,73 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroGIN", + "hidden_dim": 64, + "num_conv_layers": 4, + "pe_dim": 0, + "max_neighbours": 100, + "edge_dim": 15, + "edge_feature_names": [ + "r", + "x", + "b", + "g", + "b_fr", + "b_to", + "g_fr", + "g_to", + "rate_a", + "rate_b", + "rate_c", + "tap", + "shift", + "angmin", + "angmax" + ], + "output_heads": { + "graph": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [32, 16] + } + }, + "task_weights": [1.0], + "hetero_pooling_mode": "sum" + }, + "Variables_of_interest": { + "input_node_features": [0], + "graph_feature_names": ["objective"], + "graph_feature_dims": [1], + "node_feature_names": ["node_features"], + "node_feature_dims": [1], + "output_index": [0], + "output_dim": [1], + "output_names": ["objective"], + "type": ["graph"], + "denormalize_output": false + }, + "Training": { + "num_epoch": 2, + "batch_size": 32, + "patience": 10, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 1, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 1e-3 + } + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} diff --git a/examples/opf/opf_graph_output_homogeneous.json b/examples/opf/opf_graph_output_homogeneous.json new file mode 100644 index 000000000..a5984470a --- /dev/null +++ b/examples/opf/opf_graph_output_homogeneous.json @@ -0,0 +1,56 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "GIN", + "hidden_dim": 64, + "num_conv_layers": 4, + "pe_dim": 0, + "max_neighbours": 100, + "output_heads": { + "graph": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [32, 16] + } + }, + "task_weights": [1.0], + "edge_dim": 15, + "edge_feature_names": ["r", "x", "b", "g", "b_fr", "b_to", "g_fr", "g_to", "rate_a", "rate_b", "rate_c", "tap", "shift", "angmin", "angmax"] + }, + "Variables_of_interest": { + "input_node_features": [0], + "graph_feature_names": ["objective"], + "graph_feature_dims": [1], + "node_feature_names": ["node_features"], + "node_feature_dims": [1], + "output_index": [0], + "output_dim": [1], + "output_names": ["objective"], + "type": ["graph"], + "denormalize_output": false + }, + "Training": { + "num_epoch": 2, + "batch_size": 32, + "patience": 10, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 1, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 1e-3 + } + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} diff --git a/examples/opf/opf_heterogeneous_heat.json b/examples/opf/opf_heterogeneous_heat.json new file mode 100644 index 000000000..94b8db1c8 --- /dev/null +++ b/examples/opf/opf_heterogeneous_heat.json @@ -0,0 +1,52 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 64, + "num_conv_layers": 4, + "pe_dim": 0, + "hetero_attention_heads": 4, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "output_heads": { + "graph": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [32, 16] + } + }, + "task_weights": [1.0], + "hetero_pooling_mode": "sum" + }, + "Variables_of_interest": { + "input_node_features": [0], + "graph_feature_names": ["objective"], + "graph_feature_dims": [1], + "node_feature_names": ["node_features"], + "node_feature_dims": [1], + "output_index": [0], + "output_dim": [1], + "type": ["graph"], + "denormalize_output": false + }, + "Training": { + "num_epoch": 2, + "batch_size": 32, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 1e-3 + } + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} diff --git a/examples/opf/opf_heterogeneous_hgt.json b/examples/opf/opf_heterogeneous_hgt.json new file mode 100644 index 000000000..f669c1eed --- /dev/null +++ b/examples/opf/opf_heterogeneous_hgt.json @@ -0,0 +1,50 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHGT", + "hidden_dim": 64, + "num_conv_layers": 4, + "pe_dim": 0, + "hetero_attention_heads": 4, + "output_heads": { + "graph": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [32, 16] + } + }, + "task_weights": [1.0], + "hetero_pooling_mode": "sum" + }, + "Variables_of_interest": { + "input_node_features": [0], + "graph_feature_names": ["objective"], + "graph_feature_dims": [1], + "node_feature_names": ["node_features"], + "node_feature_dims": [1], + "output_index": [0], + "output_dim": [1], + "type": ["graph"], + "denormalize_output": false + }, + "Training": { + "num_epoch": 2, + "batch_size": 32, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 1e-3 + } + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} diff --git a/examples/opf/opf_heterogeneous_rgat.json b/examples/opf/opf_heterogeneous_rgat.json new file mode 100644 index 000000000..0771d6d2a --- /dev/null +++ b/examples/opf/opf_heterogeneous_rgat.json @@ -0,0 +1,51 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroRGAT", + "hidden_dim": 64, + "num_conv_layers": 4, + "pe_dim": 0, + "hetero_attention_heads": 4, + "hetero_attention_negative_slope": 0.2, + "output_heads": { + "graph": { + "num_sharedlayers": 1, + "dim_sharedlayers": 64, + "num_headlayers": 2, + "dim_headlayers": [32, 16] + } + }, + "task_weights": [1.0], + "hetero_pooling_mode": "sum" + }, + "Variables_of_interest": { + "input_node_features": [0], + "graph_feature_names": ["objective"], + "graph_feature_dims": [1], + "node_feature_names": ["node_features"], + "node_feature_dims": [1], + "output_index": [0], + "output_dim": [1], + "type": ["graph"], + "denormalize_output": false + }, + "Training": { + "num_epoch": 2, + "batch_size": 32, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 1e-3 + } + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} diff --git a/examples/opf/opf_nvme_utils.py b/examples/opf/opf_nvme_utils.py new file mode 100644 index 000000000..f7572655d --- /dev/null +++ b/examples/opf/opf_nvme_utils.py @@ -0,0 +1,224 @@ +import logging +import os +import shutil +import subprocess +import tarfile +import time +from mpi4py import MPI + + +def opf_release_name(topological_perturbations: bool) -> str: + return ( + "dataset_release_1_nminusone" + if topological_perturbations + else "dataset_release_1" + ) + + +def find_nvme_root(preferred_root: str | None) -> str | None: + candidates = [] + if preferred_root: + candidates.append(preferred_root) + user = os.getenv("USER", "") + if user: + candidates.append(f"/mnt/bb/{user}") + for env_var in ["LOCAL_SCRATCH", "SLURM_TMPDIR", "TMPDIR"]: + value = os.getenv(env_var) + if value: + candidates.append(value) + + for root in candidates: + if os.path.isdir(root) and os.access(root, os.W_OK | os.X_OK): + return root + return None + + +def stage_case_to_nvme( + source_datadir: str, + case_name: str, + topological_perturbations: bool, + comm, + rank: int, + nvme_root: str | None, + serialized_targets: list[str] | None = None, +) -> str: + selected_nvme_root = find_nvme_root(nvme_root) + if selected_nvme_root is None: + if rank == 0: + logging.warning( + "NVMe requested but no writable local scratch path found; using shared datadir=%s", + source_datadir, + ) + return source_datadir + + release = opf_release_name(topological_perturbations) + src_case_dir = os.path.join(source_datadir, release, case_name) + src_raw_dir = os.path.join(src_case_dir, "raw") + if not os.path.isdir(src_case_dir): + if rank == 0: + logging.warning( + "NVMe staging skipped: source case directory not found (%s)", + src_case_dir, + ) + return source_datadir + + if not os.path.isdir(src_raw_dir): + if rank == 0: + logging.warning( + "NVMe staging skipped: source raw directory not found (%s)", + src_raw_dir, + ) + return source_datadir + + archive_names = sorted( + name + for name in os.listdir(src_raw_dir) + if name.startswith(f"{case_name}_") and name.endswith(".tar.gz") + ) + if not archive_names: + if rank == 0: + logging.warning( + "NVMe staging skipped: no OPF archives found in %s", + src_raw_dir, + ) + return source_datadir + + job_id = os.getenv("SLURM_JOB_ID", "interactive") + staged_datadir = os.path.join(selected_nvme_root, "hydragnn_opf", job_id, "dataset") + dst_case_dir = os.path.join(staged_datadir, release, case_name) + dst_raw_dir = os.path.join(dst_case_dir, "raw") + copy_marker = os.path.join(dst_case_dir, ".hydragnn_archives_copied") + extract_marker = os.path.join(dst_case_dir, ".hydragnn_archives_extracted") + + effective_refresh = False + if serialized_targets: + missing_targets = [ + rel_path + for rel_path in serialized_targets + if not os.path.exists(os.path.join(staged_datadir, rel_path)) + ] + if missing_targets: + effective_refresh = True + if rank == 0: + logging.info( + "Forcing clean NVMe restore for case=%s because serialized output is missing: %s", + case_name, + ", ".join(missing_targets), + ) + + tar_path = shutil.which("tar") + + def _extract_archive(archive_path: str, dst_dir: str): + if tar_path is not None: + subprocess.run([tar_path, "-xzf", archive_path, "-C", dst_dir], check=True) + return + with tarfile.open(archive_path, "r:gz") as tar: + tar.extractall(path=dst_dir) + + local_comm = comm.Split_type(MPI.COMM_TYPE_SHARED, rank) + local_rank = local_comm.Get_rank() + local_size = local_comm.Get_size() + local_ok = True + local_err = "" + + try: + if local_rank == 0: + if effective_refresh and os.path.isdir(dst_case_dir): + shutil.rmtree(dst_case_dir, ignore_errors=True) + os.makedirs(dst_raw_dir, exist_ok=True) + local_comm.Barrier() + + did_copy = not os.path.isfile(copy_marker) + copied_local = 0 + if did_copy: + if local_rank == 0: + t0_copy = time.perf_counter() + logging.info( + "Copying OPF archives to NVMe in parallel: src=%s dst=%s archives=%d local_ranks=%d", + src_raw_dir, + dst_raw_dir, + len(archive_names), + local_size, + ) + for archive_idx in range(local_rank, len(archive_names), local_size): + archive_name = archive_names[archive_idx] + src_archive = os.path.join(src_raw_dir, archive_name) + dst_archive = os.path.join(dst_raw_dir, archive_name) + if os.path.isfile(dst_archive) and os.path.getsize( + dst_archive + ) == os.path.getsize(src_archive): + continue + shutil.copy2(src_archive, dst_archive) + copied_local += 1 + copied_total = local_comm.allreduce(copied_local, op=MPI.SUM) + local_comm.Barrier() + if local_rank == 0: + with open(copy_marker, "w") as marker: + marker.write(f"archives={len(archive_names)}\n") + logging.info( + "OPF archive copy complete for case=%s copied=%d elapsed=%.1fs", + case_name, + copied_total, + time.perf_counter() - t0_copy, + ) + + local_comm.Barrier() + + did_extract = not os.path.isfile(extract_marker) + extracted_local = 0 + if did_extract: + if local_rank == 0: + t0_extract = time.perf_counter() + logging.info( + "Extracting OPF archives on NVMe in parallel: case=%s archives=%d local_ranks=%d", + case_name, + len(archive_names), + local_size, + ) + for archive_idx in range(local_rank, len(archive_names), local_size): + archive_name = archive_names[archive_idx] + dst_archive = os.path.join(dst_raw_dir, archive_name) + _extract_archive(dst_archive, dst_raw_dir) + extracted_local += 1 + extracted_total = local_comm.allreduce(extracted_local, op=MPI.SUM) + local_comm.Barrier() + if local_rank == 0: + with open(extract_marker, "w") as marker: + marker.write(f"archives={len(archive_names)}\n") + logging.info( + "OPF archive extraction complete for case=%s extracted=%d elapsed=%.1fs", + case_name, + extracted_total, + time.perf_counter() - t0_extract, + ) + except Exception as exc: + local_ok = False + local_err = str(exc) + + node_ok = local_comm.allreduce(1 if local_ok else 0, op=MPI.MIN) == 1 + node_err = "" + if not node_ok: + node_errs = local_comm.allgather(local_err) + node_err = next((err for err in node_errs if err), "unknown local node error") + + all_ok = comm.allreduce(1 if node_ok else 0, op=MPI.MIN) == 1 + global_err = "" + if not all_ok: + global_errs = comm.allgather(node_err) + global_err = next((err for err in global_errs if err), "unknown MPI error") + local_comm.Barrier() + local_comm.Free() + comm.Barrier() + + if not all_ok: + if rank == 0: + logging.warning( + "NVMe staging failed on at least one node (%s); falling back to shared datadir=%s", + global_err, + source_datadir, + ) + return source_datadir + + if rank == 0: + logging.info("Using node-local staged datadir: %s", staged_datadir) + return staged_datadir diff --git a/examples/opf/opf_solution_heterogeneous.json b/examples/opf/opf_solution_heterogeneous.json new file mode 100644 index 000000000..91214aba6 --- /dev/null +++ b/examples/opf/opf_solution_heterogeneous.json @@ -0,0 +1,74 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroRGAT", + "hidden_dim": 64, + "num_conv_layers": 4, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": { + "num_headlayers": 2, + "dim_headlayers": [32, 16], + "type": "mlp" + } + }, + "task_weights": [1.0], + "hetero_pooling_mode": "sum", + "node_target_type": "bus" + }, + "Variables_of_interest": { + "input_node_features": [0, 1, 2, 3], + "graph_feature_names": ["context"], + "graph_feature_dims": [1], + "node_feature_names": ["node_features"], + "node_feature_dims": [4], + "output_names": ["bus_solution"], + "output_index": [0], + "output_dim": [2], + "type": ["node"], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 10, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 1, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": true, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [2, 3], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 1e-3 + } + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} diff --git a/examples/opf/opf_solution_heterogeneous_heat.json b/examples/opf/opf_solution_heterogeneous_heat.json new file mode 100644 index 000000000..761ac49ca --- /dev/null +++ b/examples/opf/opf_solution_heterogeneous_heat.json @@ -0,0 +1,94 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 64, + "num_conv_layers": 4, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16 + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 50, + "batch_size": 32, + "patience": 10, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 1, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.001 + } + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/opf_solution_homogeneous.json b/examples/opf/opf_solution_homogeneous.json new file mode 100644 index 000000000..bdb8696f7 --- /dev/null +++ b/examples/opf/opf_solution_homogeneous.json @@ -0,0 +1,56 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "PNA", + "hidden_dim": 64, + "num_conv_layers": 4, + "pe_dim": 0, + "max_neighbours": 100, + "output_heads": { + "node": { + "num_headlayers": 2, + "dim_headlayers": [32, 16], + "type": "mlp" + } + }, + "task_weights": [1.0], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "edge_dim": 11 + }, + "Variables_of_interest": { + "input_node_features": [0, 1, 2, 3], + "graph_feature_names": ["context"], + "graph_feature_dims": [1], + "node_feature_names": ["node_features"], + "node_feature_dims": [4], + "output_names": ["bus_solution"], + "output_index": [0], + "output_dim": [2], + "type": ["node"], + "denormalize_output": false + }, + "Training": { + "num_epoch": 20, + "batch_size": 32, + "patience": 10, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 1, + "continue": 0, + "startfrom": "existing_model", + "Optimizer": { + "type": "AdamW", + "learning_rate": 1e-3 + } + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} diff --git a/examples/opf/opf_solution_utils.py b/examples/opf/opf_solution_utils.py new file mode 100644 index 000000000..3f48eeb97 --- /dev/null +++ b/examples/opf/opf_solution_utils.py @@ -0,0 +1,971 @@ +"""Shared utilities for OPF solution workflows (heterogeneous and homogeneous).""" + +import copy +import logging +import os +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch_geometric.utils import degree + + +def info(*args, logtype="info", sep=" "): + getattr(logging, logtype)(sep.join(map(str, args))) + + +class OPFDomainLoss: + """Domain-informed regularization for OPF bus-level targets. + + Feasibility penalties (all zero on any strictly feasible OPF solution): + - voltage_bound_weight : Penalty for Vm (bus_pred[:, vm_output_index]) outside [v_min, v_max]. + - angle_diff_weight : Penalty for predicted Va angle-difference outside line [theta_min, theta_max]. + - line_flow_weight : Penalty for DC-approximate branch flow (DeltaVa / x_ij) exceeding rate_a. + - line_flow_slack : Tolerance subtracted from rate_a before penalising, absorbing the + small linearisation error of the DC approximation on AC-feasible + solutions. Default 1e-4 (one decade above the ~1.3e-5 residual + observed on pglib_opf_case10000_goc ground-truth data). + + Each raw penalty is normalized by a per-term exponential moving average (EMA) + before the weight is applied. This keeps every term near unit scale and makes + the weights directly comparable to the task loss, regardless of the raw + physical magnitudes (radians, per-unit power, etc.). + - ema_momentum (default 0.1): EMA decay. Smaller = slower adaptation. + + Curriculum scheduling: domain-loss weights are ramped up gradually so the + model first converges on the task loss before physics constraints are enforced. + - warmup_epochs (default 0): epochs with zero domain-loss weight. + - ramp_epochs (default 0): epochs over which weights linearly increase + from 0 to their configured values. + Example: warmup_epochs=3, ramp_epochs=3 with num_epoch=10 means: + epochs 0-2: no domain loss, epochs 3-5: linear ramp, epochs 6-9: full weight. + + Feature-index conventions (derived from the gridopt/PyG OPFDataset schema): + bus targets : [Va (0), Vm (1)] + ac_line attrs: [theta_min(0), theta_max(1), r_from(2), r_to(3), b_sh(4), x(5), rate_a(6), ...] + transformer : [theta_min(0), theta_max(1), r(2), x(3), rate_a(4), ...] + """ + + def __init__(self, config: dict | None = None, node_target_type: str = "bus"): + cfg = copy.deepcopy(config or {}) + self.enabled = bool(cfg.get("enabled", False)) + self.node_target_type = node_target_type + self.voltage_bound_weight = float(cfg.get("voltage_bound_weight", 0.0)) + self.voltage_bound_feature_indices = cfg.get( + "voltage_bound_feature_indices", None + ) + # vm_output_index: index in bus_pred corresponding to voltage magnitude (Vm). + # Default is 1 — bus targets are [Va, Vm] in the OPFDataset schema. + self.voltage_output_index = int(cfg.get("voltage_output_index", 1)) + # va_output_index: index in bus_pred corresponding to voltage angle (Va). + self.va_output_index = int(cfg.get("va_output_index", 0)) + self.angle_diff_weight = float(cfg.get("angle_diff_weight", 0.0)) + self.line_flow_weight = float(cfg.get("line_flow_weight", 0.0)) + # line_flow_slack: a small tolerance subtracted from rate_a before the DC thermal-limit + # penalty is evaluated. It exists because the DC power-flow formula + # P_ij = (Va_i - Va_j) / x_ij + # is a linearisation of the full AC power-flow equations. Even when the OPF solver + # produces a strictly AC-feasible solution, the DC approximation introduces a residual + # of ~1e-5 p.u. (empirically measured on pglib_opf_case10000_goc ground-truth data: + # mean ~1.3e-5, max ~1.7e-5). Without a slack the penalty is non-zero on ground truth, + # which means the gradient incorrectly penalises physically correct predictions. + # The default 1e-4 is one decade above the observed noise floor — large enough to zero + # out the DC-approximation artefact but small enough to still penalise real violations. + self.line_flow_slack = float(cfg.get("line_flow_slack", 1e-4)) + # EMA state for per-term scale normalization. + self._ema_momentum = float(cfg.get("ema_momentum", 0.1)) + self._penalty_ema: dict[str, float] = {} + # Curriculum scheduling. + self.warmup_epochs = int(cfg.get("warmup_epochs", 0)) + self.ramp_epochs = int(cfg.get("ramp_epochs", 0)) + + if self.voltage_bound_feature_indices is not None: + if len(self.voltage_bound_feature_indices) != 2: + raise RuntimeError( + "DomainLoss.voltage_bound_feature_indices must be [vmin_idx, vmax_idx]." + ) + self.voltage_bound_feature_indices = tuple( + int(v) for v in self.voltage_bound_feature_indices + ) + + + def _curriculum_scale(self) -> float: + """Return a [0, 1] multiplier for domain-loss weights based on current epoch. + + Reads os.environ["HYDRAGNN_EPOCH"] set by the HydraGNN training loop each + epoch — no changes to shared training code are needed. + - epoch < warmup_epochs -> 0.0 (task-loss only) + - warmup_epochs <= epoch < warmup + ramp -> linear ramp 0.0 -> 1.0 + - epoch >= warmup + ramp_epochs -> 1.0 (full weight) + """ + if self.warmup_epochs == 0 and self.ramp_epochs == 0: + return 1.0 + try: + epoch = int(os.environ.get("HYDRAGNN_EPOCH", "0")) + except (ValueError, TypeError): + return 1.0 + if epoch < self.warmup_epochs: + return 0.0 + if self.ramp_epochs <= 0: + return 1.0 + progress = (epoch - self.warmup_epochs) / self.ramp_epochs + return float(min(progress, 1.0)) + + def _normalize(self, name: str, raw: torch.Tensor) -> torch.Tensor: + """Normalize *raw* by its EMA so that the effective scale ≈ 1.0 on average. + + On the first call the EMA is seeded with the raw value, returning 1.0 + (or near-1.0 for non-zero values). Subsequent calls use the smoothed + estimate so the normalization adapts gradually as training progresses. + """ + val = float(raw.detach()) + if name not in self._penalty_ema: + # Seed: ema = raw value, normalized output = 1.0 on first step. + self._penalty_ema[name] = max(val, 1e-8) + else: + m = self._ema_momentum + self._penalty_ema[name] = max( + m * val + (1.0 - m) * self._penalty_ema[name], 1e-8 + ) + # Floor at 1e-8 prevents division by zero when a penalty term is exactly zero + # (e.g. the constraint is already satisfied for all samples in a batch). + return raw / self._penalty_ema[name] + + def __call__(self, pred, value, head_index, data): + if not self.enabled or data is None: + return value.new_zeros(()), {} + + if self.node_target_type != "bus": + return value.new_zeros(()), {} + if not hasattr(data, "node_types") or "bus" not in data.node_types: + return value.new_zeros(()), {} + if len(pred) == 0: + return value.new_zeros(()), {} + + bus_pred = pred[0] + if bus_pred.dim() == 1: + bus_pred = bus_pred.unsqueeze(-1) + bus_true = value[head_index[0]] + if bus_true.shape != bus_pred.shape: + bus_true = bus_true.reshape_as(bus_pred) + bus_true = bus_true.to(bus_pred.device) + + total_penalty = bus_pred.new_zeros(()) + metrics = {} + curriculum = self._curriculum_scale() + metrics["opf_curriculum_scale"] = torch.tensor(curriculum) + + if curriculum == 0.0: + metrics["opf_domain_total"] = total_penalty.detach() + return total_penalty, metrics + + if ( + self.voltage_bound_weight > 0.0 + and self.voltage_bound_feature_indices is not None + and hasattr(data["bus"], "x") + ): + vmin_idx, vmax_idx = self.voltage_bound_feature_indices + bus_x = data["bus"].x + if bus_x.dim() >= 2 and bus_x.shape[1] > max(vmin_idx, vmax_idx): + lower = bus_x[:, vmin_idx].reshape(-1) + upper = bus_x[:, vmax_idx].reshape(-1) + voltage = bus_pred[:, self.voltage_output_index].reshape(-1) + # F.relu zeros out values that already satisfy the bound, so the gradient + # is zero for feasible predictions and proportional to the violation otherwise. + # Squaring gives a smooth (C1) penalty with growing gradient for larger violations. + bound_penalty = torch.mean( + F.relu(lower - voltage).pow(2) + + F.relu(voltage - upper).pow(2) + ) + total_penalty = ( + total_penalty + curriculum * self.voltage_bound_weight * self._normalize("voltage_bound", bound_penalty) + ) + metrics["opf_voltage_bound"] = bound_penalty.detach() + + # ── Angle difference limit penalty ────────────────────────────────── + # Penalise predicted Va angle-differences that violate per-line bounds. + # ac_line edge_attr: [theta_min(0), theta_max(1), ...] + # transformer edge_attr: [theta_min(0), theta_max(1), ...] + if self.angle_diff_weight > 0.0 and bus_pred.shape[-1] > self.va_output_index: + Va = bus_pred[:, self.va_output_index].reshape(-1) + for rel, rel_tag in [ + (("bus", "ac_line", "bus"), "ac"), + (("bus", "transformer", "bus"), "tr"), + ]: + if rel not in data.edge_types: + continue + ea = getattr(data[rel], "edge_attr", None) + ei = getattr(data[rel], "edge_index", None) + if ea is None or ei is None or ea.numel() == 0 or ea.shape[1] < 2: + continue + theta_min = ea[:, 0].to(Va.device) + theta_max = ea[:, 1].to(Va.device) + src, dst = ei + delta_theta = Va[src] - Va[dst] + # Same relu-squared form as voltage_bound: zero gradient inside the + # feasible region [theta_min, theta_max], growing penalty outside it. + # No slack is needed here: verified empirically that this term is exactly + # zero on OPFDataset ground-truth solutions (Va and theta bounds share units). + angdiff_p = torch.mean( + F.relu(delta_theta - theta_max).pow(2) + + F.relu(theta_min - delta_theta).pow(2) + ) + total_penalty = total_penalty + curriculum * self.angle_diff_weight * self._normalize(f"{rel_tag}_angle_diff", angdiff_p) + metrics[f"opf_{rel_tag}_angle_diff"] = angdiff_p.detach() + + # ── DC thermal limit penalty ───────────────────────────────────────── + # Penalise approximate DC branch flows that exceed the thermal limit. + # P_ij = (Va_i - Va_j) / x_ij (DC power flow approximation) + # ac_line: x = edge_attr[:,5], rate_a = edge_attr[:,6] + # transformer: x = edge_attr[:,3], rate_a = edge_attr[:,4] + if self.line_flow_weight > 0.0 and bus_pred.shape[-1] > self.va_output_index: + Va = bus_pred[:, self.va_output_index].reshape(-1) + for rel, x_idx, ra_idx, rel_tag in [ + (("bus", "ac_line", "bus"), 5, 6, "ac"), + (("bus", "transformer", "bus"), 3, 4, "tr"), + ]: + if rel not in data.edge_types: + continue + ea = getattr(data[rel], "edge_attr", None) + ei = getattr(data[rel], "edge_index", None) + if ea is None or ei is None or ea.numel() == 0 or ea.shape[1] <= max(x_idx, ra_idx): + continue + # clamp x_ij away from zero to avoid division-by-zero in the DC formula; + # 1e-6 p.u. is several orders of magnitude below any physical reactance. + x_ij = ea[:, x_idx].to(Va.device).clamp(min=1e-6) + # clamp rate_a to be non-negative; negative thermal limits are nonsensical + # and could arise from edge cases in dataset normalisation. + rate_a = ea[:, ra_idx].to(Va.device).clamp(min=0.0) + src, dst = ei + # DC power-flow approximation: P_ij ≈ (Va_i - Va_j) / x_ij [per unit]. + # This linearises the full AC formula sin(Va_i - Va_j) / x_ij and is only + # exact in the flat-voltage, small-angle limit. + P_ij = (Va[src] - Va[dst]) / x_ij + # line_flow_slack is subtracted from rate_a to absorb the residual introduced + # by the DC linearisation on AC-feasible solutions (see __init__ for details). + # Without it, ground-truth predictions would incur a spurious non-zero penalty. + flow_p = torch.mean(F.relu(P_ij.abs() - rate_a - self.line_flow_slack).pow(2)) + total_penalty = total_penalty + curriculum * self.line_flow_weight * self._normalize(f"{rel_tag}_line_flow", flow_p) + metrics[f"opf_{rel_tag}_line_flow"] = flow_p.detach() + + metrics["opf_domain_total"] = total_penalty.detach() + return total_penalty, metrics + + +class OPFEnhancedModelWrapper(torch.nn.Module): + """Compose OPF-specific auxiliary loss around an existing HydraGNN model. + + In addition to combining the task loss and domain loss, this wrapper + accumulates per-batch values during each epoch and prints a one-line + breakdown at the end of that epoch (on rank 0 only). The breakdown + shows the task-driven loss and each individual domain-loss term + separately, making it straightforward to diagnose whether the domain + penalty is interfering with the data-driven objective. + + Log format (one line appended to run.log per epoch on rank 0): + DomainBreakdown epoch=XX task=X.XXXXXXXX domain_total=X.XXXXXXXX \ + curriculum=X.XX voltage_bound=X.XXXXXXXX ac_angle_diff=X.XXXXXXXX ... + """ + + def __init__(self, original_model, domain_loss: OPFDomainLoss): + super().__init__() + self.model = original_model + self.domain_loss = domain_loss + self._last_batch = None + self.last_extra_loss_metrics = {} + # Per-epoch accumulation state. + # Keyed by metric name; values are (sum, count) pairs for computing means. + self._epoch_accum: dict[str, list[float]] = {} + self._epoch_accum_task: list[float] = [] + self._last_seen_epoch: int = -1 + + def _flush_epoch_log(self, epoch: int) -> None: + """Log the mean task-loss and domain-loss breakdown for *epoch* on rank 0. + + Called automatically at the first batch of a new epoch so the previous + epoch's accumulated statistics are written before training continues. + + Log format (one line per epoch in run.log, rank 0 only):: + + LossBreakdown epoch=XX \ + data_driven_mse=X.XXXXXXXX \ + physics_penalty_total=X.XXXXXXXX \ + curriculum_scale=X.XX \ + raw_voltage_bound=X.XXXXXXXX \ + raw_ac_angle_diff=X.XXXXXXXX \ + raw_tr_angle_diff=X.XXXXXXXX \ + raw_ac_line_flow=X.XXXXXXXX + + Field meanings: + data_driven_mse -- MSE between model predictions and OPF ground-truth + targets (the standard HydraGNN task loss, no physics). + physics_penalty_total -- weighted, EMA-normalised sum of all feasibility + penalties (voltage bound + angle diff + DC flow). + This is what is added to data_driven_mse during + back-propagation. Should stay well below + data_driven_mse for the task signal to dominate. + curriculum_scale -- ramp factor in [0, 1]; 0 during warmup, 1 at full + weight. physics_penalty_total = 0 when this is 0. + raw_* -- raw (unweighted, un-normalised) value of each + individual feasibility penalty. Zero on any strictly + feasible OPF solution; non-zero indicates the current + prediction violates that constraint. + """ + # Only log from rank 0 to avoid duplicate lines in the shared run.log. + if dist.is_initialized() and dist.get_rank() != 0: + self._epoch_accum.clear() + self._epoch_accum_task.clear() + return + if not self._epoch_accum_task: + return # nothing accumulated yet (e.g. first call before any batch) + + n = len(self._epoch_accum_task) + task_mean = sum(self._epoch_accum_task) / n + + # Map internal metric keys to self-explaining log field names. + _key_labels = { + "opf_domain_total": "physics_penalty_total", + "opf_curriculum_scale": "curriculum_scale", + "opf_voltage_bound": "raw_voltage_bound", + "opf_ac_angle_diff": "raw_ac_angle_diff", + "opf_tr_angle_diff": "raw_tr_angle_diff", + "opf_ac_line_flow": "raw_ac_line_flow", + "opf_tr_line_flow": "raw_tr_line_flow", + } + + parts = [f"epoch={epoch:02d}", f"data_driven_mse={task_mean:.8f}"] + for key in sorted(self._epoch_accum): + vals = self._epoch_accum[key] + mean_val = sum(vals) / len(vals) + label = _key_labels.get(key, key.removeprefix("opf_")) + parts.append(f"{label}={mean_val:.8f}") + + # Use print rather than logging.info so the line is always visible in + # run.log regardless of the logging level configured by HydraGNN. + print("0: LossBreakdown " + " ".join(parts), flush=True) + + # Reset accumulators for the next epoch. + self._epoch_accum.clear() + self._epoch_accum_task.clear() + + def __getattr__(self, name): + try: + return super().__getattr__(name) + except AttributeError: + return getattr(self.model, name) + + def forward(self, data): + self._last_batch = data + return self.model(data) + + def loss(self, pred, value, head_index): + total_loss, tasks_loss = self.model.loss(pred, value, head_index) + if self._last_batch is None: + info( + "[OPFEnhancedModelWrapper] loss() called before forward(); " + "domain penalty will be zero for this batch.", + logtype="warning", + ) + extra_loss, extra_metrics = self.domain_loss( + pred, + value, + head_index, + self._last_batch, + ) + self.last_extra_loss_metrics = extra_metrics + + # ── Per-epoch accumulation ─────────────────────────────────────────── + # Detect epoch transitions using HYDRAGNN_EPOCH (set by the core training + # loop). On each new epoch, flush the previous epoch's accumulated stats + # to logging.info so they appear in run.log alongside the Epoch: XX line. + try: + current_epoch = int(os.environ.get("HYDRAGNN_EPOCH", "-1")) + except (ValueError, TypeError): + current_epoch = -1 + + if current_epoch != self._last_seen_epoch and self._last_seen_epoch >= 0: + # Epoch boundary: flush accumulated stats for the completed epoch. + self._flush_epoch_log(self._last_seen_epoch) + self._last_seen_epoch = current_epoch + + # Accumulate task loss (total_loss is the data-driven term before domain is added). + self._epoch_accum_task.append(float(total_loss.detach())) + # Accumulate each domain metric (raw, un-normalized values for interpretability). + for key, val in extra_metrics.items(): + self._epoch_accum.setdefault(key, []).append(float(val)) + + return total_loss + extra_loss, tasks_loss + + +def build_solution_target(data, node_target_type: str): + """Extract the solution target tensor for the given node type.""" + if hasattr(data, "node_types") and node_target_type in data.node_types: + node_store = data[node_target_type] + if not hasattr(node_store, "y") or node_store.y is None: + raise RuntimeError( + f"No targets found for node type '{node_target_type}' in OPF sample." + ) + return node_store.y.to(torch.float32) + + if hasattr(data, "_node_type_names") and hasattr(data, "node_type"): + if node_target_type not in data._node_type_names: + raise RuntimeError( + f"Node type '{node_target_type}' not found in OPF sample." + ) + type_index = data._node_type_names.index(node_target_type) + if not hasattr(data, "y") or data.y is None: + raise RuntimeError( + f"No homogeneous targets found for node type '{node_target_type}'." + ) + mask = data.node_type == type_index + return data.y[mask].to(torch.float32) + + raise RuntimeError(f"Node type '{node_target_type}' not found in OPF sample.") + + +def ensure_node_y_loc(data): + if not hasattr(data, "y") or data.y is None: + raise RuntimeError("Missing node targets (data.y) for OPF sample.") + if data.y.dim() == 1: + data.y = data.y.unsqueeze(-1) + num_nodes = int(data.y.shape[0]) + target_dim = int(data.y.shape[1]) + data.y_num_nodes = torch.tensor( + [num_nodes], dtype=torch.int64, device=data.y.device + ) + data.y_loc = torch.tensor( + [[0, num_nodes * target_dim]], + dtype=torch.int64, + device=data.y.device, + ) + + +def resolve_node_target_type(data, requested: str) -> str: + if hasattr(data, "node_types"): + if requested in data.node_types: + return requested + if hasattr(data, "_node_type_names") and requested in data._node_type_names: + idx = data._node_type_names.index(requested) + if idx < len(data.node_types): + return data.node_types[idx] + raise RuntimeError( + f"Requested node_target_type '{requested}' not found in data. " + f"Available node types: {list(data.node_types)}." + ) + if hasattr(data, "_node_type_names") and requested in data._node_type_names: + return requested + raise RuntimeError( + f"Cannot resolve node_target_type '{requested}': data has no node_types." + ) + + +def _as_edge_feature(value, num_edges: int, device): + if value is None: + return None + if not isinstance(value, torch.Tensor): + try: + value = torch.as_tensor(value) + except Exception: + return None + if value.numel() == 0: + return None + if value.dim() == 0: + return None + if value.dim() == 1: + if int(value.shape[0]) != int(num_edges): + return None + value = value.view(-1, 1) + elif value.dim() >= 2: + if int(value.shape[0]) != int(num_edges): + return None + value = value.reshape(num_edges, -1) + if value.dtype not in (torch.float16, torch.float32, torch.float64): + value = value.to(torch.float32) + return value.to(device=device, dtype=torch.float32) + + +def resolve_edge_feature_schema( + configured_feature_names=None, + configured_edge_dim=None, +): + if configured_feature_names is None or len(configured_feature_names) == 0: + raise RuntimeError( + "edge_feature_names must be explicitly provided in the config. " + "No implicit defaults are used." + ) + schema = [str(name) for name in configured_feature_names if str(name).strip()] + if not schema: + raise RuntimeError("edge_feature_names contains only empty/whitespace entries.") + if configured_edge_dim is not None: + edge_dim = int(configured_edge_dim) + if edge_dim != len(schema): + raise RuntimeError( + f"edge_dim={edge_dim} does not match the number of " + f"edge_feature_names ({len(schema)}). They must be equal." + ) + return tuple(schema) + + +def validate_voi_node_features(config: dict, node_target_type: str | None = None): + """Validate that node feature config is fully specified. Crash on anything missing.""" + nn_config = config.get("NeuralNetwork") + if nn_config is None: + raise RuntimeError("Config is missing 'NeuralNetwork' section.") + var_config = nn_config.get("Variables_of_interest") + if var_config is None: + raise RuntimeError("Config is missing 'NeuralNetwork.Variables_of_interest'.") + + input_node_features = var_config.get("input_node_features") + if not isinstance(input_node_features, list) or len(input_node_features) == 0: + raise RuntimeError( + "'input_node_features' must be an explicit non-empty list in the config." + ) + + node_feature_dims = var_config.get("node_feature_dims") + if not isinstance(node_feature_dims, list) or len(node_feature_dims) == 0: + raise RuntimeError( + "'node_feature_dims' must be an explicit non-empty list in the config." + ) + + if "node_feature_names" not in var_config: + raise RuntimeError( + "'node_feature_names' must be explicitly provided in the config." + ) + + return config + + +def compute_pna_deg_for_hetero_dataset(dataset, verbosity: int = 2): + from hydragnn.utils.print.print_utils import iterate_tqdm + + num_samples = len(dataset) + if dist.is_initialized(): + rank = dist.get_rank() + world_size = dist.get_world_size() + start = (num_samples * rank) // world_size + end = (num_samples * (rank + 1)) // world_size + else: + start = 0 + end = num_samples + + local_indices = range(start, end) + + max_deg_local = 0 + for idx in iterate_tqdm(local_indices, verbosity, desc="HeteroPNA degree max"): + data = dataset[idx] + data_h = data.to_homogeneous(add_node_type=True, add_edge_type=True) + d = degree(data_h.edge_index[1], num_nodes=data_h.num_nodes, dtype=torch.long) + if d.numel() > 0: + max_deg_local = max(max_deg_local, int(d.max().item())) + + if dist.is_initialized(): + reduce_device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + max_deg_tensor = torch.tensor( + [max_deg_local], dtype=torch.long, device=reduce_device + ) + dist.all_reduce(max_deg_tensor, op=dist.ReduceOp.MAX) + max_deg = int(max_deg_tensor.item()) + else: + max_deg = max_deg_local + + deg_local = torch.zeros(max_deg + 1, dtype=torch.long) + for idx in iterate_tqdm(local_indices, verbosity, desc="HeteroPNA degree bincount"): + data = dataset[idx] + data_h = data.to_homogeneous(add_node_type=True, add_edge_type=True) + d = degree(data_h.edge_index[1], num_nodes=data_h.num_nodes, dtype=torch.long) + deg_local += torch.bincount(d, minlength=deg_local.numel()) + + if dist.is_initialized(): + reduce_device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + deg_tensor = deg_local.to(device=reduce_device) + dist.all_reduce(deg_tensor, op=dist.ReduceOp.SUM) + deg = deg_tensor.cpu() + else: + deg = deg_local + + return deg.tolist() + + +def _assemble_edge_attr_hetero(data, edge_dim_dict): + """Heterogeneous route: keep per-edge-type native widths. + + Edge types whose relation name appears in *edge_dim_dict* must carry a + pre-assembled ``edge_attr`` tensor with the declared width. Edge types + absent from the dict are treated as featureless — any stale ``edge_attr`` + is removed so that ``data.edge_attr_dict`` only contains featured types. + + Returns ``(data, edge_dim_dict)`` unchanged. + """ + for edge_type in data.edge_types: + _, rel, _ = edge_type + edge_store = data[edge_type] + edge_index = getattr(edge_store, "edge_index", None) + if not isinstance(edge_index, torch.Tensor): + continue + if edge_index.dim() != 2 or edge_index.size(0) != 2: + continue + + expected_dim = edge_dim_dict.get(rel) + existing = getattr(edge_store, "edge_attr", None) + + if expected_dim is None: + # Featureless — remove any edge_attr so it stays out of + # data.edge_attr_dict during training. + if existing is not None: + try: + delattr(edge_store, "edge_attr") + except AttributeError: + pass + continue + + if not isinstance(existing, torch.Tensor) or existing.dim() != 2: + raise RuntimeError( + f"Edge type {edge_type} (rel={rel}) expects edge_attr with " + f"{expected_dim} columns but found no valid 2-D tensor." + ) + if existing.size(1) != expected_dim: + raise RuntimeError( + f"Edge type {edge_type} (rel={rel}) has edge_attr width " + f"{existing.size(1)}, expected {expected_dim} from edge_dim config." + ) + + return data, edge_dim_dict + + +def assemble_edge_attr(data, edge_dim, feature_schema=None): + """One-time assembly during preprocessing. + + *edge_dim* determines the route: + + * **int** — *homogeneous* route. Every edge type is zero-padded (or + assembled from named columns via *feature_schema*) to a uniform width + equal to *edge_dim*. + * **dict** — *heterogeneous* route. Keys are relation names (the middle + element of an edge-type triple); values are the expected widths of + pre-assembled ``edge_attr`` tensors. Edge types absent from the dict + are treated as featureless. + + Returns ``(data, edge_dim)``. + """ + if not hasattr(data, "edge_types"): + return data, edge_dim + + if isinstance(edge_dim, dict): + return _assemble_edge_attr_hetero(data, edge_dim) + + target_dim = int(edge_dim) + if target_dim <= 0: + raise RuntimeError("int edge_dim must be positive.") + + schema = None + if feature_schema is not None: + schema = tuple(str(n) for n in feature_schema if str(n).strip()) + if not schema: + schema = None + + for edge_type in data.edge_types: + edge_store = data[edge_type] + edge_index = getattr(edge_store, "edge_index", None) + if not isinstance(edge_index, torch.Tensor): + continue + if edge_index.dim() != 2 or edge_index.size(0) != 2: + continue + num_edges = int(edge_index.size(1)) + device = edge_index.device + + # Already assembled — accept as-is or zero-pad to target_dim. + existing = getattr(edge_store, "edge_attr", None) + if ( + isinstance(existing, torch.Tensor) + and existing.dim() == 2 + and existing.size(0) == num_edges + ): + w = existing.size(1) + if w == target_dim: + continue # exact match + if w < target_dim: + pad = torch.zeros( + num_edges, target_dim - w, device=device, dtype=existing.dtype + ) + data[edge_type].edge_attr = torch.cat([existing, pad], dim=1) + continue + raise RuntimeError( + f"edge_attr for {edge_type} has {w} columns, exceeding edge_dim={target_dim}." + ) + + # Try named-column assembly if a schema was provided. + if schema is not None: + has_any = any( + getattr(edge_store, name, None) is not None for name in schema + ) + if not has_any and existing is None: + data[edge_type].edge_attr = torch.zeros( + num_edges, target_dim, device=device, dtype=torch.float32 + ) + continue + + cols = [] + for attr_name in schema: + col = _as_edge_feature( + getattr(edge_store, attr_name, None), num_edges, device + ) + if col is None: + raise RuntimeError( + f"Missing or invalid edge attribute '{attr_name}' " + f"for edge type {edge_type}." + ) + if int(col.shape[1]) != 1: + raise RuntimeError( + f"Edge attribute '{attr_name}' for edge type {edge_type} has " + f"{int(col.shape[1])} columns; expected exactly 1." + ) + cols.append(col) + + data[edge_type].edge_attr = torch.cat(cols, dim=1).contiguous() + + for attr_name in schema: + try: + delattr(edge_store, attr_name) + except AttributeError: + pass + continue + + # No schema and no existing tensor — zero-fill. + if existing is None: + data[edge_type].edge_attr = torch.zeros( + num_edges, target_dim, device=device, dtype=torch.float32 + ) + + return data, target_dim + + +def _validate_edge_attr_hetero(data, edge_dim_dict): + """Check per-edge-type widths for the heterogeneous route.""" + for edge_type in data.edge_types: + _, rel, _ = edge_type + edge_store = data[edge_type] + edge_index = getattr(edge_store, "edge_index", None) + if not isinstance(edge_index, torch.Tensor): + continue + if edge_index.dim() != 2 or edge_index.size(0) != 2: + continue + num_edges = int(edge_index.size(1)) + + expected_dim = edge_dim_dict.get(rel) + edge_attr = getattr(edge_store, "edge_attr", None) + + if expected_dim is None: + # Featureless — must NOT have edge_attr. + if isinstance(edge_attr, torch.Tensor): + raise RuntimeError( + f"Featureless edge type {edge_type} (rel={rel}) should not " + f"have edge_attr, but found tensor with shape {list(edge_attr.shape)}." + ) + continue + + if not isinstance(edge_attr, torch.Tensor): + raise RuntimeError( + f"Edge type {edge_type} (rel={rel}) is missing edge_attr; " + f"expected width {expected_dim}." + ) + if edge_attr.dim() != 2: + raise RuntimeError( + f"edge_attr for edge type {edge_type} has {edge_attr.dim()} " + f"dimensions; expected 2." + ) + if edge_attr.size(0) != num_edges: + raise RuntimeError( + f"edge_attr row count mismatch for edge type {edge_type}: " + f"got {edge_attr.size(0)}, expected {num_edges}." + ) + if edge_attr.size(1) != expected_dim: + raise RuntimeError( + f"edge_attr dim mismatch for edge type {edge_type} (rel={rel}): " + f"got {edge_attr.size(1)}, expected {expected_dim}." + ) + + return data + + +def validate_edge_attr(data, edge_dim): + """Validate that every edge type carries properly shaped ``edge_attr``. + + *edge_dim* can be: + + * **int** — every edge type must have ``edge_attr`` with that many columns + (featureless types that have no ``edge_attr`` are silently skipped). + * **dict** — per-relation-name widths; featureless types (absent from the + dict) must NOT carry ``edge_attr``. + """ + if not hasattr(data, "edge_types"): + return data + + if isinstance(edge_dim, dict): + return _validate_edge_attr_hetero(data, edge_dim) + + target_dim = int(edge_dim) + + for edge_type in data.edge_types: + edge_store = data[edge_type] + edge_index = getattr(edge_store, "edge_index", None) + if not isinstance(edge_index, torch.Tensor): + continue + if edge_index.dim() != 2 or edge_index.size(0) != 2: + continue + num_edges = int(edge_index.size(1)) + + edge_attr = getattr(edge_store, "edge_attr", None) + if not isinstance(edge_attr, torch.Tensor): + continue + if edge_attr.dim() != 2: + raise RuntimeError( + f"edge_attr for edge type {edge_type} has " + f"{edge_attr.dim()} dimensions; expected 2." + ) + if edge_attr.size(0) != num_edges: + raise RuntimeError( + f"edge_attr row count mismatch for edge type {edge_type}: " + f"got {edge_attr.size(0)}, expected {num_edges}." + ) + if edge_attr.size(1) != target_dim: + raise RuntimeError( + f"edge_attr dim mismatch for edge type {edge_type}: " + f"got {edge_attr.size(1)}, expected {target_dim}." + ) + + return data + + +class HeteroFromHomogeneousDataset: + """Wraps an ADIOS-loaded homogeneous dataset, converting each sample to + heterogeneous and validating ``edge_attr`` shape. + """ + + def __init__(self, base, edge_dim: int): + self.base = base + self.edge_dim = edge_dim + + def __len__(self): + return len(self.base) + + def __getitem__(self, idx): + data = self.base[idx] + hetero = data.to_heterogeneous() + if hasattr(data, "y"): + hetero.y = data.y + if hasattr(data, "graph_attr"): + hetero.graph_attr = data.graph_attr + validate_edge_attr(hetero, self.edge_dim) + return hetero + + +class EdgeAttrDatasetAdapter: + """Validates ``edge_attr`` on every access — no assembly, just shape check.""" + + def __init__(self, base, edge_dim: int): + self.base = base + self.edge_dim = edge_dim + + def __len__(self): + return len(self.base) + + def __getitem__(self, idx): + data = self.base[idx] + validate_edge_attr(data, self.edge_dim) + return data + + def __getattr__(self, name): + return getattr(self.base, name) + + +class NodeTargetDatasetAdapter: + def __init__(self, base, node_target_type: str, edge_dim: int): + self.base = base + self.node_target_type = node_target_type + self.edge_dim = edge_dim + + def __len__(self): + return len(self.base) + + def __getitem__(self, idx): + data = self.base[idx] + validate_edge_attr(data, self.edge_dim) + if ( + not hasattr(data, "node_types") + or self.node_target_type not in data.node_types + ): + raise RuntimeError( + f"Node type '{self.node_target_type}' not found in OPF sample." + ) + if ( + not hasattr(data[self.node_target_type], "y") + or data[self.node_target_type].y is None + ): + raise RuntimeError( + f"No targets found for node type '{self.node_target_type}' in OPF sample." + ) + data.y = data[self.node_target_type].y + ensure_node_y_loc(data) + return data + + def __getattr__(self, name): + return getattr(self.base, name) + + +class NodeBatchAdapter: + def __init__(self, loader, node_target_type: str, edge_dim: int): + self.loader = loader + self.node_target_type = node_target_type + self.edge_dim = edge_dim + self.dataset = loader.dataset + self.sampler = getattr(loader, "sampler", None) + + def __iter__(self): + for data in self.loader: + validate_edge_attr(data, self.edge_dim) + if ( + not hasattr(data, "node_types") + or self.node_target_type not in data.node_types + ): + raise RuntimeError( + f"Node type '{self.node_target_type}' not found in OPF sample." + ) + + if not hasattr(data, "batch"): + node_store = data[self.node_target_type] + if hasattr(node_store, "batch"): + data.batch = node_store.batch + elif ( + hasattr(data, "batch_dict") + and self.node_target_type in data.batch_dict + ): + data.batch = data.batch_dict[self.node_target_type] + else: + raise RuntimeError( + f"Cannot find batch vector for node type " + f"'{self.node_target_type}' in batched OPF data." + ) + + if ( + not hasattr(data[self.node_target_type], "y") + or data[self.node_target_type].y is None + ): + raise RuntimeError( + f"No targets found for node type '{self.node_target_type}' in OPF sample." + ) + data.y = data[self.node_target_type].y + ensure_node_y_loc(data) + yield data + + def __len__(self): + return len(self.loader) + + def __getattr__(self, name): + return getattr(self.loader, name) diff --git a/examples/opf/plot_hpo_combined.py b/examples/opf/plot_hpo_combined.py new file mode 100644 index 000000000..9646eb66a --- /dev/null +++ b/examples/opf/plot_hpo_combined.py @@ -0,0 +1,212 @@ +#!/usr/bin/env python3 +""" +Generate combined HPO validation loss plot from multiple HPO runs, +correctly pairing each CSV source with its DeepHyper output directory. +""" + +import csv +import os +import re +from collections import defaultdict + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.ticker as mticker +import numpy as np + +TYPE_COLORS = { + "HeteroHEAT": "#d62728", + "HeteroHGT": "#1f77b4", + "HeteroPNA": "#2ca02c", + "HeteroSAGE": "#ff7f0e", + "HeteroGAT": "#9467bd", + "HeteroRGAT": "#8c564b", + "HeteroGIN": "#e377c2", +} +TYPE_MARKERS = { + "HeteroHEAT": "s", + "HeteroHGT": "o", + "HeteroPNA": "D", + "HeteroSAGE": "^", + "HeteroGAT": "v", + "HeteroRGAT": "X", + "HeteroGIN": "P", +} + + +def extract_epoch_losses(dh_dir, trial_id): + path = os.path.join(dh_dir, "output-0.{}.txt".format(trial_id)) + if not os.path.exists(path): + return [] + losses = [] + with open(path) as f: + for line in f: + m = re.match(r"Val loss:\s+([\d.eE+-]+)", line.strip()) + if m: + val = float(m.group(1)) + if np.isfinite(val): + losses.append(val) + return losses + + +def load_run(csv_path, dh_dir): + """Load one HPO run, correctly pairing CSV rows with their dh_dir.""" + trials = [] + n_failed = 0 + failed_types = defaultdict(int) + with open(csv_path) as f: + for r in csv.DictReader(f): + if r["objective"] == "F": + n_failed += 1 + failed_types[r["p:mpnn_type"]] += 1 + continue + trial_id = r["job_id"] + mpnn = r["p:mpnn_type"] + losses = extract_epoch_losses(dh_dir, trial_id) + if not losses: + continue + hparams = { + "hidden_dim": int(r["p:hidden_dim"]), + "num_conv_layers": int(r["p:num_conv_layers"]), + "learning_rate": float(r["p:learning_rate"]), + } + trials.append((mpnn, trial_id, hparams, losses)) + return trials, n_failed, dict(failed_types) + + +def main(): + # Define the two HPO runs: (csv_path, dh_dir) + runs = [ + ("opf_hpo-4249563/results.csv", "deephyper-opf-hpo-4249563"), + ("opf_hpo_HeteroPNA-4252004/results.csv", "deephyper-opf-hpo-4252004"), + ] + + by_type = defaultdict(list) + total_failed = 0 + all_failed_types = defaultdict(int) + + for csv_path, dh_dir in runs: + trials, n_failed, failed_types = load_run(csv_path, dh_dir) + total_failed += n_failed + for ft, fc in failed_types.items(): + all_failed_types[ft] += fc + for mpnn, trial_id, hparams, losses in trials: + by_type[mpnn].append((trial_id, hparams, losses)) + print("Loaded {} successful trials from {} (dh_dir={})".format( + len(trials), csv_path, dh_dir)) + + total = sum(len(v) for v in by_type.values()) + print("\nTotal: {} successful, {} failed".format(total, total_failed)) + + # Print summary + print("\n{:>4} {:>14} {:>6} {:>6} {:>10} {:>13}".format( + "Rank", "Model Type", "Hidden", "Layers", "LR", "Best Val Loss")) + print("-" * 60) + all_trials = [] + for mpnn, trials in by_type.items(): + for trial_id, hparams, losses in trials: + all_trials.append({ + "mpnn": mpnn, "trial": trial_id, + "hidden": hparams["hidden_dim"], + "layers": hparams["num_conv_layers"], + "lr": hparams["learning_rate"], + "best_loss": min(losses), + }) + all_trials.sort(key=lambda x: x["best_loss"]) + for i, t in enumerate(all_trials): + print("{:>4} {:>14} {:>6} {:>6} {:>10.6f} {:>13.6f}".format( + i+1, t["mpnn"], t["hidden"], t["layers"], t["lr"], t["best_loss"])) + + # Per-type summary + print("\n{:>14} {:>6} {:>10} {:>10} {:>10}".format( + "Model Type", "Trials", "Best", "Mean", "Worst")) + print("-" * 54) + type_stats = {} + for mpnn, trials in by_type.items(): + losses_all = [min(l) for _, _, l in trials] + type_stats[mpnn] = { + "count": len(losses_all), + "best": min(losses_all), + "worst": max(losses_all), + "mean": sum(losses_all) / len(losses_all), + } + for mpnn in sorted(type_stats, key=lambda t: type_stats[t]["best"]): + s = type_stats[mpnn] + print("{:>14} {:>6} {:>10.6f} {:>10.6f} {:>10.6f}".format( + mpnn, s["count"], s["best"], s["mean"], s["worst"])) + + # --- Plot --- + fig, ax = plt.subplots(figsize=(12, 7)) + type_order = sorted( + by_type.keys(), key=lambda t: min(min(l) for _, _, l in by_type[t])) + + legend_handles = [] + legend_labels = [] + max_epochs = 0 + overall_best_loss = float("inf") + overall_best_info = {} + + for mpnn in type_order: + trials = by_type[mpnn] + color = TYPE_COLORS.get(mpnn, "#333333") + marker = TYPE_MARKERS.get(mpnn, "o") + trials.sort(key=lambda t: min(t[2])) + + for idx, (trial_id, hparams, losses) in enumerate(trials): + epochs = list(range(1, len(losses) + 1)) + max_epochs = max(max_epochs, len(losses)) + alpha = 0.9 if idx == 0 else 0.35 + lw = 2.5 if idx == 0 else 1.0 + line = ax.semilogy( + epochs, losses, + color=color, marker=marker, + markersize=6 if idx == 0 else 4, + linewidth=lw, alpha=alpha, + markeredgewidth=0.5, markeredgecolor="white") + if idx == 0: + best_loss = min(losses) + legend_handles.append(line[0]) + legend_labels.append("{} (best: {:.5f})".format(mpnn, best_loss)) + if best_loss < overall_best_loss: + overall_best_loss = best_loss + best_epoch = losses.index(best_loss) + 1 + overall_best_info = { + "mpnn": mpnn, "trial": trial_id, + "loss": best_loss, "epoch": best_epoch} + + ax.set_xlabel("Epoch", fontsize=13, fontweight="bold") + ax.set_ylabel("Validation Loss", fontsize=13, fontweight="bold") + ax.set_title( + "HPO Validation Loss Curves by Model Architecture\n" + "({} trials, Jobs 4249563 + 4252004, 128 Frontier nodes)".format(total), + fontsize=14, fontweight="bold") + ax.set_xticks(range(1, max_epochs + 1)) + ax.set_xlim(0.5, max_epochs + 0.5) + ax.grid(True, which="both", alpha=0.3, linestyle="--") + ax.yaxis.set_minor_formatter(mticker.NullFormatter()) + + ax.legend(legend_handles, legend_labels, + loc="upper right", fontsize=10, framealpha=0.9, + title="Model Type (best trial val loss)", title_fontsize=10) + + if overall_best_info: + info = overall_best_info + ax.annotate( + "Best overall:\n{}, trial {}\nVal loss = {:.5f}".format( + info["mpnn"], info["trial"], info["loss"]), + xy=(info["epoch"], info["loss"]), + xytext=(max(1, info["epoch"] - 3), info["loss"] * 2.5), + fontsize=9, + arrowprops=dict(arrowstyle="->", color="gray", lw=1.2), + bbox=dict(boxstyle="round,pad=0.3", + facecolor="lightyellow", edgecolor="gray")) + + plt.tight_layout() + plot_path = "hpo_validation_loss_curves.png" + plt.savefig(plot_path, dpi=200, bbox_inches="tight") + print("\nPlot saved to {}".format(plot_path)) + + +if __name__ == "__main__": + main() diff --git a/examples/opf/plot_hpo_results.py b/examples/opf/plot_hpo_results.py new file mode 100644 index 000000000..e22bc6f19 --- /dev/null +++ b/examples/opf/plot_hpo_results.py @@ -0,0 +1,659 @@ +#!/usr/bin/env python3 +""" +Analyze DeepHyper HPO results for HydraGNN OPF heterogeneous graph training. + +Produces: + 1. A semilogy plot of validation loss curves per epoch, colored by model type. + 2. A console summary table of all trials and per-model-type statistics. + 3. (Optional) Appends an HPO results section to the OPF explanation docx. + +Usage examples: + # Minimal: just specify the job ID (assumes standard directory layout) + python plot_hpo_results.py --job_id 4249563 + + # Explicit paths (when directories were renamed or moved) + python plot_hpo_results.py \\ + --csv_path opf_hpo-4249563/results.csv \\ + --dh_dir deephyper-opf-hpo-4249563 \\ + --plot_path hpo_validation_loss_curves.png + + # Also update the docx + python plot_hpo_results.py --job_id 4249563 \\ + --docx_path OPF_Heterogeneous_Graph_Explanation.docx + + # Combine multiple HPO runs + python plot_hpo_results.py \\ + --csv_path opf_hpo-4249563/results.csv opf_hpo-4249999/results.csv \\ + --dh_dir deephyper-opf-hpo-4249563 deephyper-opf-hpo-4249999 +""" + +import argparse +import csv +import os +import re +import sys +from collections import defaultdict + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import matplotlib.ticker as mticker +import numpy as np + +# ─── Visual style ──────────────────────────────────────────────────────────── +TYPE_COLORS = { + "HeteroHEAT": "#d62728", # red + "HeteroHGT": "#1f77b4", # blue + "HeteroPNA": "#2ca02c", # green + "HeteroSAGE": "#ff7f0e", # orange + "HeteroGAT": "#9467bd", # purple + "HeteroRGAT": "#8c564b", # brown + "HeteroGIN": "#e377c2", # pink +} + +TYPE_MARKERS = { + "HeteroHEAT": "s", + "HeteroHGT": "o", + "HeteroPNA": "D", + "HeteroSAGE": "^", + "HeteroGAT": "v", + "HeteroRGAT": "X", + "HeteroGIN": "P", +} + + +# ─── Data loading ──────────────────────────────────────────────────────────── +def load_results(csv_paths): + """Load and merge results from one or more DeepHyper CSV files.""" + all_rows = [] + for path in csv_paths: + with open(path) as f: + rows = list(csv.DictReader(f)) + print(f" Loaded {len(rows)} trials from {path}") + all_rows.extend(rows) + return all_rows + + +def extract_epoch_losses(dh_dirs, trial_id): + """Read output-0.{trial_id}.txt from the first matching dh_dir.""" + for dh_dir in dh_dirs: + path = os.path.join(dh_dir, f"output-0.{trial_id}.txt") + if not os.path.exists(path): + continue + losses = [] + with open(path) as f: + for line in f: + m = re.match(r"Val loss:\s+([\d.eE+-]+)", line.strip()) + if m: + val = float(m.group(1)) + if np.isfinite(val): + losses.append(val) + if losses: + return losses + return [] + + +def build_trial_data(rows, dh_dirs): + """Returns dict: mpnn_type -> list of (trial_id, hyperparams, [epoch_losses])""" + by_type = defaultdict(list) + n_failed = 0 + failed_types = defaultdict(int) + for r in rows: + if r["objective"] == "F": + n_failed += 1 + failed_types[r["p:mpnn_type"]] += 1 + continue + trial_id = r["job_id"] + mpnn = r["p:mpnn_type"] + losses = extract_epoch_losses(dh_dirs, trial_id) + if not losses: + continue + hparams = { + "hidden_dim": int(r["p:hidden_dim"]), + "num_conv_layers": int(r["p:num_conv_layers"]), + "learning_rate": float(r["p:learning_rate"]), + } + by_type[mpnn].append((trial_id, hparams, losses)) + return by_type, n_failed, dict(failed_types) + + +# ─── Plotting ──────────────────────────────────────────────────────────────── +def create_plot(by_type, plot_path, title_extra=""): + """Create semilogy plot of val loss curves, colored by model type.""" + fig, ax = plt.subplots(figsize=(12, 7)) + + type_order = sorted( + by_type.keys(), key=lambda t: min(min(l) for _, _, l in by_type[t]) + ) + + legend_handles = [] + legend_labels = [] + max_epochs = 0 + overall_best_loss = float("inf") + overall_best_info = {} + + for mpnn in type_order: + trials = by_type[mpnn] + color = TYPE_COLORS.get(mpnn, "#333333") + marker = TYPE_MARKERS.get(mpnn, "o") + trials.sort(key=lambda t: min(t[2])) + + for idx, (trial_id, hparams, losses) in enumerate(trials): + epochs = list(range(1, len(losses) + 1)) + max_epochs = max(max_epochs, len(losses)) + alpha = 0.9 if idx == 0 else 0.35 + lw = 2.5 if idx == 0 else 1.0 + + line = ax.semilogy( + epochs, + losses, + color=color, + marker=marker, + markersize=6 if idx == 0 else 4, + linewidth=lw, + alpha=alpha, + markeredgewidth=0.5, + markeredgecolor="white", + ) + + if idx == 0: + best_loss = min(losses) + legend_handles.append(line[0]) + legend_labels.append(f"{mpnn} (best: {best_loss:.5f})") + if best_loss < overall_best_loss: + overall_best_loss = best_loss + best_epoch = losses.index(best_loss) + 1 + overall_best_info = { + "mpnn": mpnn, + "trial": trial_id, + "loss": best_loss, + "epoch": best_epoch, + } + + total_trials = sum(len(v) for v in by_type.values()) + ax.set_xlabel("Epoch", fontsize=13, fontweight="bold") + ax.set_ylabel("Validation Loss", fontsize=13, fontweight="bold") + + title = f"HPO Validation Loss Curves by Model Architecture\n({total_trials} trials" + if title_extra: + title += f", {title_extra}" + title += ")" + ax.set_title(title, fontsize=14, fontweight="bold") + + ax.set_xticks(range(1, max_epochs + 1)) + ax.set_xlim(0.5, max_epochs + 0.5) + ax.grid(True, which="both", alpha=0.3, linestyle="--") + ax.yaxis.set_minor_formatter(mticker.NullFormatter()) + + ax.legend( + legend_handles, + legend_labels, + loc="upper right", + fontsize=10, + framealpha=0.9, + title="Model Type (best trial val loss)", + title_fontsize=10, + ) + + if overall_best_info: + info = overall_best_info + ax.annotate( + f"Best overall:\n{info['mpnn']}, trial {info['trial']}\n" + f"Val loss = {info['loss']:.5f}", + xy=(info["epoch"], info["loss"]), + xytext=(max(1, info["epoch"] - 3), info["loss"] * 2.5), + fontsize=9, + arrowprops=dict(arrowstyle="->", color="gray", lw=1.2), + bbox=dict( + boxstyle="round,pad=0.3", + facecolor="lightyellow", + edgecolor="gray", + ), + ) + + plt.tight_layout() + plt.savefig(plot_path, dpi=200, bbox_inches="tight") + print(f"Plot saved to {plot_path}") + return plot_path + + +# ─── Console summary ───────────────────────────────────────────────────────── +def print_summary(by_type, n_failed, failed_types): + """Print formatted trial summary to console.""" + total = sum(len(v) for v in by_type.values()) + print(f"\n{'='*75}") + print(f" HPO Results: {total} successful trials, {n_failed} failed") + if failed_types: + parts = [f"{t}: {n}" for t, n in sorted(failed_types.items())] + print(f" Failed by type: {', '.join(parts)}") + print(f"{'='*75}\n") + + # All trials ranked + all_trials = [] + for mpnn, trials in by_type.items(): + for trial_id, hparams, losses in trials: + all_trials.append( + { + "trial": trial_id, + "mpnn": mpnn, + "hidden": hparams["hidden_dim"], + "layers": hparams["num_conv_layers"], + "lr": hparams["learning_rate"], + "best_loss": min(losses), + } + ) + all_trials.sort(key=lambda x: x["best_loss"]) + + print(f"{'Rank':>4} {'Trial':>5} {'Model Type':>14} {'Hidden':>6} {'Layers':>6} {'LR':>10} {'Best Val Loss':>13}") + print("-" * 68) + for i, t in enumerate(all_trials): + print( + f"{i+1:>4} {t['trial']:>5} {t['mpnn']:>14} {t['hidden']:>6} " + f"{t['layers']:>6} {t['lr']:>10.6f} {t['best_loss']:>13.6f}" + ) + + # Per-type summary + print(f"\n{'--- Summary by Model Type ---':^68}") + print(f"{'Model Type':>14} {'Trials':>6} {'Best':>10} {'Mean':>10} {'Worst':>10}") + print("-" * 54) + type_stats = {} + for mpnn, trials in by_type.items(): + losses_all = [min(l) for _, _, l in trials] + type_stats[mpnn] = { + "count": len(losses_all), + "best": min(losses_all), + "worst": max(losses_all), + "mean": sum(losses_all) / len(losses_all), + } + for mpnn in sorted(type_stats, key=lambda t: type_stats[t]["best"]): + s = type_stats[mpnn] + print( + f"{mpnn:>14} {s['count']:>6} {s['best']:>10.6f} " + f"{s['mean']:>10.6f} {s['worst']:>10.6f}" + ) + print() + return all_trials, type_stats + + +# ─── Document update ───────────────────────────────────────────────────────── +def update_docx(by_type, plot_path, docx_path, n_failed, failed_types): + """Append HPO results section to the OPF explanation docx.""" + from docx import Document + from docx.enum.text import WD_ALIGN_PARAGRAPH + from docx.shared import Inches + + doc = Document(docx_path) + + # Check if section 15.4 already exists — avoid duplicating + for p in doc.paragraphs: + if "15.4" in p.text and "HPO Results" in p.text: + print(f"WARNING: Section '15.4 HPO Results' already exists in {docx_path}.") + print(" Skipping docx update to avoid duplication.") + print(" Delete the existing section first if you want to regenerate it.") + return + + total = sum(len(v) for v in by_type.values()) + failed_detail = "" + if failed_types: + parts = [f"{t}" for t in sorted(failed_types)] + failed_detail = ( + f" (all {', '.join(parts)} configurations, which crashed due to NCCL " + "communication errors)" + ) + + def add_heading(text, level=2): + p = doc.add_paragraph() + p.style = doc.styles[f"Heading {level}"] + p.text = text + return p + + def add_normal(text): + p = doc.add_paragraph() + p.style = doc.styles["Normal"] + p.text = text + return p + + def add_bullet(text): + p = doc.add_paragraph() + p.style = doc.styles["List Bullet"] + p.text = text + return p + + # ── Section 15.4: HPO Results ── + add_heading("15.4 HPO Results", level=2) + add_normal( + f"The HPO search completed {total + n_failed} trials, of which {total} " + f"completed successfully and {n_failed} failed{failed_detail}. " + "Despite any partial completion, the results provide strong evidence for " + "which architectures and hyperparameter regions perform best on the AC-OPF problem." + ) + + # ── Results Table ── + add_heading("15.4.1 Trial Results Summary", level=2) + add_normal( + "The following table shows the top 10 HPO trials ranked by best validation loss " + "(minimum across all epochs). The Bayesian optimizer successfully identified " + "the top-performing architectures and concentrated its search in the most " + "promising hyperparameter regions." + ) + + all_success = [] + for mpnn, trials in by_type.items(): + for trial_id, hparams, losses in trials: + all_success.append( + { + "trial": trial_id, + "mpnn": mpnn, + "hidden": hparams["hidden_dim"], + "layers": hparams["num_conv_layers"], + "lr": hparams["learning_rate"], + "best_loss": min(losses), + } + ) + all_success.sort(key=lambda x: x["best_loss"]) + + n_show = min(10, len(all_success)) + table = doc.add_table(rows=n_show + 1, cols=6) + table.style = "Light Shading Accent 1" + headers = [ + "Rank", "Model Type", "Hidden Dim", + "Conv Layers", "Learning Rate", "Best Val Loss", + ] + for j, h in enumerate(headers): + cell = table.rows[0].cells[j] + cell.text = h + for run in cell.paragraphs[0].runs: + run.bold = True + + for i, trial in enumerate(all_success[:n_show]): + row = table.rows[i + 1] + row.cells[0].text = str(i + 1) + row.cells[1].text = trial["mpnn"] + row.cells[2].text = str(trial["hidden"]) + row.cells[3].text = str(trial["layers"]) + row.cells[4].text = f"{trial['lr']:.6f}" + row.cells[5].text = f"{trial['best_loss']:.6f}" + + add_normal("") + + # ── Summary by Model Type ── + add_heading("15.4.2 Performance by Model Architecture", level=2) + add_normal( + "Aggregating results by model architecture reveals clear performance tiers. " + "The table below summarizes the validation loss statistics for each model type." + ) + + type_stats = {} + for mpnn, trials in by_type.items(): + losses_all = [min(l) for _, _, l in trials] + type_stats[mpnn] = { + "count": len(losses_all), + "best": min(losses_all), + "worst": max(losses_all), + "mean": sum(losses_all) / len(losses_all), + } + type_order = sorted(type_stats.keys(), key=lambda t: type_stats[t]["best"]) + + table2 = doc.add_table(rows=len(type_order) + 1, cols=5) + table2.style = "Light Shading Accent 1" + for j, h in enumerate( + ["Model Type", "Trials", "Best Val Loss", "Mean Val Loss", "Worst Val Loss"] + ): + cell = table2.rows[0].cells[j] + cell.text = h + for run in cell.paragraphs[0].runs: + run.bold = True + + for i, mpnn in enumerate(type_order): + s = type_stats[mpnn] + row = table2.rows[i + 1] + row.cells[0].text = mpnn + row.cells[1].text = str(s["count"]) + row.cells[2].text = f"{s['best']:.6f}" + row.cells[3].text = f"{s['mean']:.6f}" + row.cells[4].text = f"{s['worst']:.6f}" + + add_normal("") + + # ── Training curves figure ── + add_heading("15.4.3 Validation Loss Training Curves", level=2) + add_normal( + "The following figure shows the validation loss curves (semilogy scale) for all " + f"{total} successful HPO trials, colored by model architecture. For each model " + "type, the best-performing trial is drawn with a bold line, while other trials " + "of the same type appear as thinner, more transparent lines." + ) + + doc.add_picture(plot_path, width=Inches(6.0)) + doc.paragraphs[-1].alignment = WD_ALIGN_PARAGRAPH.CENTER + + # Build dynamic caption from actual best/worst types + hex_to_name = { + "#d62728": "red", "#1f77b4": "blue", "#2ca02c": "green", + "#ff7f0e": "orange", "#9467bd": "purple", "#8c564b": "brown", + "#e377c2": "pink", "#333333": "gray", + } + best_two = type_order[:2] + worst_two = type_order[-2:] if len(type_order) >= 4 else type_order[-1:] + color_map = {t: TYPE_COLORS.get(t, "#333333") for t in type_order} + best_desc = " and ".join( + f"{t} ({hex_to_name.get(color_map[t], 'colored')})" for t in best_two + ) + worst_desc = " and ".join( + f"{t} ({hex_to_name.get(color_map[t], 'colored')})" for t in worst_two + ) + add_normal( + f"Figure: Validation loss vs. epoch for all successful HPO trials. " + f"{best_desc} consistently achieve the lowest validation losses, while " + f"{worst_desc} remain significantly worse." + ) + + add_normal("") + + # ── Analysis ── + add_heading("15.4.4 Analysis of HPO Results", level=2) + add_normal( + "The HPO results provide several key insights that refine the preliminary " + "findings from Section 14:" + ) + + best_trial = all_success[0] + add_bullet( + f"{best_trial['mpnn']} emerges as the top architecture. With a best validation " + f"loss of {best_trial['best_loss']:.5f} (trial {best_trial['trial']}: " + f"hidden_dim={best_trial['hidden']}, {best_trial['layers']} conv layers, " + f"lr={best_trial['lr']:.5f}), it outperforms all other architectures." + ) + + if len(type_order) >= 2: + second_type = type_order[1] + second_best = type_stats[second_type] + add_bullet( + f"{second_type} is a strong second. With a best validation loss of " + f"{second_best['best']:.5f} across {second_best['count']} trials " + f"(mean {second_best['mean']:.5f}), it demonstrates consistent performance." + ) + + # Common hyperparameter patterns in top 5 + top5 = all_success[:min(5, len(all_success))] + layers_set = set(t["layers"] for t in top5) + if len(layers_set) == 1: + common_layers = top5[0]["layers"] + add_bullet( + f"Optimal hyperparameter region. The top {len(top5)} trials all use " + f"{common_layers} convolutional layers with hidden dimensions between " + f"{min(t['hidden'] for t in top5)} and {max(t['hidden'] for t in top5)} " + f"and learning rates between {min(t['lr'] for t in top5):.4f} and " + f"{max(t['lr'] for t in top5):.4f}." + ) + else: + add_bullet( + f"Optimal hyperparameter region. The top {len(top5)} trials use " + f"{min(t['layers'] for t in top5)}-{max(t['layers'] for t in top5)} " + f"convolutional layers with hidden dimensions between " + f"{min(t['hidden'] for t in top5)} and {max(t['hidden'] for t in top5)}." + ) + + if failed_types: + for ft, fc in sorted(failed_types.items(), key=lambda x: -x[1]): + add_bullet( + f"{ft}: all {fc} trial(s) failed, suggesting scalability or " + "stability issues in the distributed training setting." + ) + + for mpnn in type_order[-2:]: + if type_stats[mpnn]["best"] > 2 * type_stats[type_order[0]]["best"]: + add_bullet( + f"{mpnn} underperforms significantly with a best validation loss of " + f"{type_stats[mpnn]['best']:.5f} — " + f"{type_stats[mpnn]['best']/type_stats[type_order[0]]['best']:.1f}x " + f"worse than {type_order[0]}." + ) + + add_normal("") + + # ── Recommendation ── + add_heading("15.4.5 Revised Model Recommendation", level=2) + second = type_order[1] if len(type_order) >= 2 else "N/A" + add_normal( + f"Based on the HPO results, {best_trial['mpnn']} is recommended as the primary " + f"architecture for the AC-OPF problem, with {second} as a strong alternative. " + f"The optimal configuration is: {best_trial['mpnn']} with " + f"{best_trial['layers']} convolutional layers, hidden dimension " + f"~{best_trial['hidden']}, and learning rate ~{best_trial['lr']:.4f}." + ) + + doc.save(docx_path) + print(f"Document updated: {docx_path}") + + +# ─── CLI ───────────────────────────────────────────────────────────────────── +def parse_args(): + parser = argparse.ArgumentParser( + description="Analyze DeepHyper HPO results for HydraGNN OPF training.", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + # Quick run with just a job ID (uses standard directory naming): + %(prog)s --job_id 4249563 + + # Explicit paths: + %(prog)s --csv_path opf_hpo-4249563/results.csv \\ + --dh_dir deephyper-opf-hpo-4249563 + + # Merge multiple HPO runs: + %(prog)s --csv_path opf_hpo-4249563/results.csv opf_hpo-4250000/results.csv \\ + --dh_dir deephyper-opf-hpo-4249563 deephyper-opf-hpo-4250000 + + # Also update the docx: + %(prog)s --job_id 4249563 --docx_path OPF_Heterogeneous_Graph_Explanation.docx +""", + ) + + parser.add_argument( + "--job_id", + type=str, + default=None, + help="Slurm job ID. If provided, csv_path and dh_dir are auto-derived " + "as 'opf_hpo-{job_id}/results.csv' and 'deephyper-opf-hpo-{job_id}/'.", + ) + parser.add_argument( + "--csv_path", + nargs="+", + default=None, + help="Path(s) to DeepHyper results.csv file(s). Overrides --job_id.", + ) + parser.add_argument( + "--dh_dir", + nargs="+", + default=None, + help="Path(s) to DeepHyper output directories containing output-0.*.txt " + "files with per-epoch val losses. Overrides --job_id.", + ) + parser.add_argument( + "--plot_path", + type=str, + default="hpo_validation_loss_curves.png", + help="Output path for the semilogy plot image (default: hpo_validation_loss_curves.png).", + ) + parser.add_argument( + "--title_extra", + type=str, + default="", + help="Extra text appended to the plot title (e.g., '128 Frontier nodes').", + ) + parser.add_argument( + "--docx_path", + type=str, + default=None, + help="Path to the OPF_Heterogeneous_Graph_Explanation.docx. " + "If provided, HPO results section 15.4 will be appended.", + ) + parser.add_argument( + "--no_plot", + action="store_true", + help="Skip plot generation (print summary only).", + ) + return parser.parse_args() + + +def main(): + args = parse_args() + + # Resolve paths + if args.csv_path: + csv_paths = args.csv_path + elif args.job_id: + csv_paths = [f"opf_hpo-{args.job_id}/results.csv"] + else: + print("ERROR: Provide either --job_id or --csv_path.", file=sys.stderr) + sys.exit(1) + + if args.dh_dir: + dh_dirs = args.dh_dir + elif args.job_id: + dh_dirs = [f"deephyper-opf-hpo-{args.job_id}"] + else: + print("ERROR: Provide either --job_id or --dh_dir.", file=sys.stderr) + sys.exit(1) + + # Validate inputs + for p in csv_paths: + if not os.path.exists(p): + print(f"ERROR: CSV file not found: {p}", file=sys.stderr) + sys.exit(1) + for d in dh_dirs: + if not os.path.isdir(d): + print(f"WARNING: DeepHyper dir not found: {d} (epoch curves will be missing)") + + # Load and build data + print("Loading results...") + rows = load_results(csv_paths) + by_type, n_failed, failed_types = build_trial_data(rows, dh_dirs) + + total = sum(len(v) for v in by_type.values()) + if total == 0: + print("ERROR: No successful trials with epoch data found.", file=sys.stderr) + sys.exit(1) + + # Console summary + print_summary(by_type, n_failed, failed_types) + + # Plot + if not args.no_plot: + title_extra = args.title_extra + if not title_extra and args.job_id: + title_extra = f"Job {args.job_id}" + create_plot(by_type, args.plot_path, title_extra=title_extra) + + # Docx update + if args.docx_path: + if not os.path.exists(args.docx_path): + print(f"ERROR: Docx file not found: {args.docx_path}", file=sys.stderr) + sys.exit(1) + update_docx(by_type, args.plot_path, args.docx_path, n_failed, failed_types) + + +if __name__ == "__main__": + main() diff --git a/examples/opf/pretrained_models/HeteroHEAT_best/HeteroHEAT_best.pk b/examples/opf/pretrained_models/HeteroHEAT_best/HeteroHEAT_best.pk new file mode 100644 index 000000000..c3ea00454 Binary files /dev/null and b/examples/opf/pretrained_models/HeteroHEAT_best/HeteroHEAT_best.pk differ diff --git a/examples/opf/pretrained_models/HeteroHEAT_best/config.json b/examples/opf/pretrained_models/HeteroHEAT_best/config.json new file mode 100644 index 000000000..03df57030 --- /dev/null +++ b/examples/opf/pretrained_models/HeteroHEAT_best/config.json @@ -0,0 +1,149 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroHEAT", + "hidden_dim": 232, + "num_conv_layers": 6, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "avg_num_neighbors": null, + "radius": null, + "radial_type": null, + "distance_transform": null, + "num_gaussians": null, + "num_filters": null, + "envelope_exponent": null, + "num_after_skip": null, + "num_before_skip": null, + "basis_emb_size": null, + "int_emb_size": null, + "out_emb_size": null, + "num_radial": null, + "num_spherical": null, + "correlation": null, + "max_ell": null, + "node_max_ell": null, + "enable_interatomic_potential": false, + "equivariance": null, + "freeze_conv_layers": false, + "initial_bias": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 10, + "batch_size": 32, + "patience": 10, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 1, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": true, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.009923017058624745 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/pretrained_models/HeteroHEAT_best/opf_hpo_4466179_0.11_epoch_9.pk b/examples/opf/pretrained_models/HeteroHEAT_best/opf_hpo_4466179_0.11_epoch_9.pk new file mode 100644 index 000000000..c3ea00454 Binary files /dev/null and b/examples/opf/pretrained_models/HeteroHEAT_best/opf_hpo_4466179_0.11_epoch_9.pk differ diff --git a/examples/opf/pretrained_models/HeteroSAGE_best/HeteroSAGE_best.pk b/examples/opf/pretrained_models/HeteroSAGE_best/HeteroSAGE_best.pk new file mode 100644 index 000000000..4e64be1c4 Binary files /dev/null and b/examples/opf/pretrained_models/HeteroSAGE_best/HeteroSAGE_best.pk differ diff --git a/examples/opf/pretrained_models/HeteroSAGE_best/config.json b/examples/opf/pretrained_models/HeteroSAGE_best/config.json new file mode 100644 index 000000000..726e0ed0e --- /dev/null +++ b/examples/opf/pretrained_models/HeteroSAGE_best/config.json @@ -0,0 +1,149 @@ +{ + "Verbosity": { + "level": 2 + }, + "NeuralNetwork": { + "Architecture": { + "mpnn_type": "HeteroSAGE", + "hidden_dim": 141, + "num_conv_layers": 5, + "pe_dim": 0, + "max_neighbours": 100, + "hetero_attention_heads": 4, + "edge_dim": { + "ac_line": 9, + "transformer": 11 + }, + "node_input_dims": { + "bus": 4, + "generator": 11, + "load": 2, + "shunt": 2 + }, + "output_heads": { + "node": [ + { + "type": "branch-0", + "architecture": { + "num_headlayers": 2, + "dim_headlayers": [ + 32, + 16 + ], + "type": "mlp" + } + } + ] + }, + "task_weights": [ + 1.0 + ], + "hetero_pooling_mode": "sum", + "node_target_type": "bus", + "global_attn_engine": null, + "global_attn_type": null, + "global_attn_heads": 0, + "output_dim": [ + 2 + ], + "output_type": [ + "node" + ], + "num_nodes": 31, + "input_dim": 4, + "pna_deg": null, + "avg_num_neighbors": null, + "radius": null, + "radial_type": null, + "distance_transform": null, + "num_gaussians": null, + "num_filters": null, + "envelope_exponent": null, + "num_after_skip": null, + "num_before_skip": null, + "basis_emb_size": null, + "int_emb_size": null, + "out_emb_size": null, + "num_radial": null, + "num_spherical": null, + "correlation": null, + "max_ell": null, + "node_max_ell": null, + "enable_interatomic_potential": false, + "equivariance": null, + "freeze_conv_layers": false, + "initial_bias": null, + "activation_function": "relu", + "hetero_attention_negative_slope": 0.2, + "hetero_edge_type_emb_dim": 16, + "hetero_edge_attr_emb_dim": 16, + "SyncBatchNorm": false + }, + "Variables_of_interest": { + "input_node_features": [ + 0, + 1, + 2, + 3 + ], + "graph_feature_names": [ + "context" + ], + "graph_feature_dims": [ + 1 + ], + "node_feature_names": [ + "node_features" + ], + "node_feature_dims": [ + 4 + ], + "output_names": [ + "bus_solution" + ], + "output_index": [ + 0 + ], + "output_dim": [ + 2 + ], + "type": [ + "node" + ], + "denormalize_output": false + }, + "Training": { + "num_epoch": 10, + "batch_size": 32, + "patience": 10, + "early_stopping": true, + "Checkpoint": true, + "checkpoint_warmup": 1, + "continue": 0, + "startfrom": "existing_model", + "DomainLoss": { + "enabled": true, + "smoothness_weight": 0.001, + "transformer_smoothness_weight": 0.001, + "voltage_bound_weight": 0.01, + "voltage_bound_feature_indices": [ + 2, + 3 + ], + "voltage_output_index": -1 + }, + "Optimizer": { + "type": "AdamW", + "learning_rate": 0.0038913632520334433 + }, + "conv_checkpointing": false, + "loss_function_type": "mse", + "precision": "fp32" + } + }, + "Visualization": { + "plot_init_solution": false, + "plot_hist_solution": false, + "create_plots": false + } +} \ No newline at end of file diff --git a/examples/opf/pretrained_models/HeteroSAGE_best/opf_hpo_4466182_0.13_epoch_9.pk b/examples/opf/pretrained_models/HeteroSAGE_best/opf_hpo_4466182_0.13_epoch_9.pk new file mode 100644 index 000000000..4e64be1c4 Binary files /dev/null and b/examples/opf/pretrained_models/HeteroSAGE_best/opf_hpo_4466182_0.13_epoch_9.pk differ diff --git a/examples/opf/train_opf_graph_output_heterogeneous.py b/examples/opf/train_opf_graph_output_heterogeneous.py new file mode 100644 index 000000000..ead8d00cb --- /dev/null +++ b/examples/opf/train_opf_graph_output_heterogeneous.py @@ -0,0 +1,740 @@ +import os +import json +import logging +import argparse +import shutil +import subprocess +import sys +from mpi4py import MPI + +import torch +import torch.distributed as dist +from torch_geometric.datasets import OPFDataset +import torch_geometric.datasets.opf as tg_opf +from __init__ import data_ops +from opf_nvme_utils import stage_case_to_nvme +from opf_solution_utils import ( + EdgeAttrDatasetAdapter, + HeteroFromHomogeneousDataset, + assemble_edge_attr, + compute_pna_deg_for_hetero_dataset, + info, + resolve_edge_feature_schema, + validate_voi_node_features, +) + + +def _to_jsonable(obj): + if isinstance(obj, torch.Tensor): + return obj.item() if obj.numel() == 1 else obj.tolist() + if isinstance(obj, dict): + return {k: _to_jsonable(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_to_jsonable(v) for v in obj] + return obj + + +def _patch_fast_tar_extraction(): + tar_path = shutil.which("tar") + if tar_path is None: + return + + original_extract_tar = tg_opf.extract_tar + + def _fast_extract_tar(path: str, folder: str, mode: str = "r:gz", log: bool = True): + if log: + print(f"Extracting {path}", file=sys.stderr) + try: + try: + subprocess.run( + [ + tar_path, + "--checkpoint=1000", + "--checkpoint-action=dot", + "-xzf", + path, + "-C", + folder, + ], + check=True, + ) + if log: + print("", file=sys.stderr) + return + except Exception: + subprocess.run( + [tar_path, "-xzf", path, "-C", folder], + check=True, + ) + except Exception: + original_extract_tar(path, folder, mode=mode, log=log) + + tg_opf.extract_tar = _fast_extract_tar + + +import hydragnn +import time + + +def _diag(msg: str): + if os.getenv("HYDRAGNN_DIAG") != "1": + return + rank = dist.get_rank() if dist.is_initialized() else 0 + rank_filter = os.getenv("HYDRAGNN_DIAG_RANK") + if rank_filter is not None: + try: + if int(rank_filter) != int(rank): + return + except ValueError: + pass + now = time.perf_counter() + print(f"[diag][rank {rank}][{now:.3f}] {msg}", flush=True) + + +from hydragnn.utils.datasets.pickledataset import ( + SimplePickleWriter, + SimplePickleDataset, +) +from hydragnn.utils.distributed import nsplit +from hydragnn.utils.model import print_model +from hydragnn.utils.input_config_parsing.config_utils import update_config + +try: + from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset +except ImportError: + AdiosWriter = None + AdiosDataset = None + + +def _prepare_sample( + data, + to_homogeneous: bool, + edge_dim=None, + edge_feature_schema=None, +): + data.y = data.objective.view(1, 1).to(torch.float32) + data.graph_attr = data.x.view(1, -1).to(torch.float32) + data, _ = assemble_edge_attr( + data, edge_dim=edge_dim, feature_schema=edge_feature_schema + ) + if not to_homogeneous: + return data + data_h = data.to_homogeneous( + node_attrs=["x"], + edge_attrs=["edge_attr"], + add_node_type=True, + add_edge_type=True, + ) + data_h.y = data.y + data_h.graph_attr = data.graph_attr + return data_h + + +def _load_split(root, split, case_name, num_groups, topological_perturbations): + if dist.is_initialized(): + if dist.get_rank() == 0: + info( + f"Loading OPF split: case={case_name} split={split} groups={num_groups}" + ) + else: + info(f"Loading OPF split: case={case_name} split={split} groups={num_groups}") + + def _construct(force_reload: bool = False): + return OPFDataset( + root=root, + split=split, + case_name=case_name, + num_groups=num_groups, + topological_perturbations=topological_perturbations, + force_reload=force_reload, + ) + + try: + return _construct(force_reload=False) + except Exception as exc: + msg = str(exc) + recoverable = isinstance(exc, EOFError) or any( + token in msg + for token in ( + "PytorchStreamReader failed reading file", + "PytorchStreamReader failed reading zip archive", + "failed finding central directory", + "Cannot use ``weights_only=True`` with files saved in the legacy .tar format", + "Weights only load failed", + "Unsupported operand 80", + ) + ) + if not recoverable: + raise + + if dist.is_initialized() and dist.get_world_size() > 1: + rank = dist.get_rank() + info( + f"Corrupted processed cache detected for case={case_name} split={split}; rank {rank} rebuilding with force_reload=True" + ) + dataset = _construct(force_reload=True) + MPI.COMM_WORLD.Barrier() + return dataset + + info( + f"Corrupted processed cache detected for case={case_name} split={split}; rebuilding with force_reload=True" + ) + return _construct(force_reload=True) + + +def _prime_processed_splits_on_rank0( + root, + case_name, + num_groups, + topological_perturbations, + rank, + comm, +): + if rank == 0: + for split in ("train", "val", "test"): + _load_split( + root, + split, + case_name, + num_groups, + topological_perturbations, + ) + comm.Barrier() + + +def _subset_for_rank(dataset, rank, world_size): + rx = list(nsplit(range(len(dataset)), world_size))[rank] + return [dataset[i] for i in range(rx.start, rx.stop)] + + +def _log_phase_time(comm, rank, label: str, elapsed_local: float): + elapsed_max = comm.allreduce(float(elapsed_local), op=MPI.MAX) + elapsed_sum = comm.allreduce(float(elapsed_local), op=MPI.SUM) + elapsed_avg = elapsed_sum / max(1, comm.Get_size()) + if rank == 0: + info(f"Timing {label}: max={elapsed_max:.2f}s avg={elapsed_avg:.2f}s") + + +def _resolve_preonly_case_names(args, datadir): + if not args.preonly_case_names: + return [args.case_name] + + case_names = [c.strip() for c in args.preonly_case_names if c.strip()] + if len(case_names) == 1 and case_names[0].lower() == "all": + discovered = data_ops.discover_cases(datadir, args.topological_perturbations) + if discovered: + return discovered + return [args.case_name] + return case_names + + +if __name__ == "__main__": + _patch_fast_tar_extraction() + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "--inputfile", type=str, default="opf_graph_output_heterogeneous.json" + ) + parser.add_argument("--data_root", type=str, default="dataset") + parser.add_argument( + "--case_name", + type=str, + default="pglib_opf_case14_ieee", + ) + parser.add_argument( + "--num_groups", + type=str, + default="1", + help="Number of groups or 'all'", + ) + parser.add_argument( + "--num_groups_max", + type=int, + default=20, + help="Fallback/probe cap when --num_groups all and none on disk", + ) + parser.add_argument( + "--no_num_groups_probe", + action="store_false", + dest="num_groups_probe", + help="Disable probing remote storage when --num_groups all and none on disk", + ) + parser.set_defaults(num_groups_probe=True) + parser.add_argument("--topological_perturbations", action="store_true") + parser.add_argument("--preonly", action="store_true", help="preprocess only") + parser.add_argument( + "--preonly_case_names", + nargs="+", + default=None, + help="Case names to preprocess when --preonly is set; supports 'all'", + ) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--num_epoch", type=int, default=None) + parser.add_argument("--modelname", type=str, default="OPF_Hetero") + parser.add_argument( + "--nvme", + action="store_true", + help="Stage selected OPF case onto node-local NVMe/scratch if available", + ) + + group = parser.add_mutually_exclusive_group() + group.add_argument("--adios", action="store_const", dest="format", const="adios") + group.add_argument("--pickle", action="store_const", dest="format", const="pickle") + parser.set_defaults(format="pickle") + + args = parser.parse_args() + + dirpwd = os.path.dirname(os.path.abspath(__file__)) + datadir = os.path.join(dirpwd, args.data_root) + input_filename = os.path.join(dirpwd, args.inputfile) + + with open(input_filename, "r") as f: + config = json.load(f) + + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + raw_edge_dim = arch_config.get("edge_dim") + if isinstance(raw_edge_dim, dict): + edge_dim = {str(k): int(v) for k, v in raw_edge_dim.items()} + edge_feature_schema = None + elif raw_edge_dim is not None: + edge_dim = int(raw_edge_dim) + names = arch_config.get("edge_feature_names") + if names: + edge_feature_schema = resolve_edge_feature_schema(names, edge_dim) + else: + edge_feature_schema = None + else: + raise RuntimeError("edge_dim must be specified in config.") + arch_config["edge_dim"] = edge_dim + + if args.batch_size is not None: + config["NeuralNetwork"]["Training"]["batch_size"] = args.batch_size + if args.num_epoch is not None: + config["NeuralNetwork"]["Training"]["num_epoch"] = args.num_epoch + + comm_size, rank = hydragnn.utils.distributed.setup_ddp() + comm = MPI.COMM_WORLD + + logging.basicConfig( + level=logging.INFO, + format=f"%(levelname)s (rank {rank}): %(message)s", + datefmt="%H:%M:%S", + ) + + log_name = args.modelname + hydragnn.utils.print.setup_log(log_name) + writer = hydragnn.utils.model.get_summary_writer(log_name) + + requested_num_groups = data_ops.parse_num_groups(args.num_groups) + adios_training_only = args.format == "adios" and not args.preonly + + if args.preonly: + case_names = _resolve_preonly_case_names(args, datadir) + trainset, valset, testset = [], [], [] + output_datadir = datadir + store_homogeneous = args.format == "adios" + + for case_name in case_names: + t_case = time.perf_counter() + num_groups = data_ops.resolve_num_groups( + requested_num_groups, + datadir, + case_name, + args.topological_perturbations, + args.num_groups_max, + args.num_groups_probe, + rank, + comm, + ) + + t_download = time.perf_counter() + data_ops.ensure_opf_downloaded( + datadir, + case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} phase=download_extract groups={num_groups}", + time.perf_counter() - t_download, + ) + + case_datadir = datadir + if args.nvme: + serialized_target = ( + f"{args.modelname}.bp" + if args.format == "adios" + else f"{args.modelname}.pickle" + ) + t_stage = time.perf_counter() + case_datadir = stage_case_to_nvme( + datadir, + case_name, + args.topological_perturbations, + comm, + rank, + None, + serialized_targets=[serialized_target], + ) + _log_phase_time( + comm, + rank, + f"case={case_name} phase=nvme_stage groups={num_groups}", + time.perf_counter() - t_stage, + ) + if args.format != "adios": + output_datadir = case_datadir + + _prime_processed_splits_on_rank0( + case_datadir, + case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + + t_load = time.perf_counter() + train_raw = _load_split( + case_datadir, + "train", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=train phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_load = time.perf_counter() + val_raw = _load_split( + case_datadir, + "val", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=val phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_load = time.perf_counter() + test_raw = _load_split( + case_datadir, + "test", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=test phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_pre = time.perf_counter() + train_subset = _subset_for_rank(train_raw, rank, comm_size) + trainset.extend( + _prepare_sample( + d, + store_homogeneous, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + for d in train_subset + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=train phase=preprocess local_samples={len(train_subset)}", + time.perf_counter() - t_pre, + ) + + t_pre = time.perf_counter() + val_subset = _subset_for_rank(val_raw, rank, comm_size) + valset.extend( + _prepare_sample( + d, + store_homogeneous, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + for d in val_subset + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=val phase=preprocess local_samples={len(val_subset)}", + time.perf_counter() - t_pre, + ) + + t_pre = time.perf_counter() + test_subset = _subset_for_rank(test_raw, rank, comm_size) + testset.extend( + _prepare_sample( + d, + store_homogeneous, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + for d in test_subset + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=test phase=preprocess local_samples={len(test_subset)}", + time.perf_counter() - t_pre, + ) + + _log_phase_time( + comm, + rank, + f"case={case_name} phase=prepare_total groups={num_groups}", + time.perf_counter() - t_case, + ) + + info( + f"Local split sizes: train={len(trainset)}, val={len(valset)}, test={len(testset)}" + ) + + if args.format == "adios": + t_write = time.perf_counter() + if AdiosWriter is None: + raise RuntimeError("adios2 is not available in this environment.") + fname = os.path.join(output_datadir, f"{args.modelname}.bp") + if rank == 0 and os.path.exists(fname): + if os.path.isdir(fname): + shutil.rmtree(fname, ignore_errors=True) + else: + os.remove(fname) + comm.Barrier() + adwriter = AdiosWriter(fname, comm) + adwriter.add("trainset", trainset) + adwriter.add("valset", valset) + adwriter.add("testset", testset) + adwriter.save() + _log_phase_time( + comm, + rank, + f"phase=serialize format=adios model={args.modelname}", + time.perf_counter() - t_write, + ) + else: + t_write = time.perf_counter() + basedir = os.path.join(output_datadir, f"{args.modelname}.pickle") + SimplePickleWriter(trainset, basedir, "trainset", use_subdir=True) + SimplePickleWriter(valset, basedir, "valset", use_subdir=True) + SimplePickleWriter(testset, basedir, "testset", use_subdir=True) + _log_phase_time( + comm, + rank, + f"phase=serialize format=pickle model={args.modelname}", + time.perf_counter() - t_write, + ) + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() + raise SystemExit(0) + + if not adios_training_only: + num_groups = data_ops.resolve_num_groups( + requested_num_groups, + datadir, + args.case_name, + args.topological_perturbations, + args.num_groups_max, + args.num_groups_probe, + rank, + comm, + ) + + data_ops.ensure_opf_downloaded( + datadir, + args.case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + + if args.nvme: + serialized_target = ( + f"{args.modelname}.bp" + if args.format == "adios" + else f"{args.modelname}.pickle" + ) + datadir = stage_case_to_nvme( + datadir, + args.case_name, + args.topological_perturbations, + comm, + rank, + None, + serialized_targets=[serialized_target], + ) + + _prime_processed_splits_on_rank0( + datadir, + args.case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + + info("Loading OPF splits...") + train_raw = _load_split( + datadir, + "train", + args.case_name, + num_groups, + args.topological_perturbations, + ) + val_raw = _load_split( + datadir, + "val", + args.case_name, + num_groups, + args.topological_perturbations, + ) + test_raw = _load_split( + datadir, + "test", + args.case_name, + num_groups, + args.topological_perturbations, + ) + + # Validate var_config from config — no auto-fill from data + var_config = config["NeuralNetwork"]["Variables_of_interest"] + validate_voi_node_features(config) + if ( + not isinstance(var_config.get("graph_feature_dims"), list) + or len(var_config["graph_feature_dims"]) == 0 + ): + raise RuntimeError( + "'graph_feature_dims' must be an explicit non-empty list in the config." + ) + + if args.format == "adios": + if AdiosDataset is None: + raise RuntimeError("adios2 is not available in this environment.") + fname = os.path.join(datadir, f"{args.modelname}.bp") + if adios_training_only and not os.path.isdir(fname): + raise RuntimeError( + f"Expected preprocessed ADIOS dataset at '{fname}' for training-only mode. " + "Run with --preonly --adios first." + ) + train_base = AdiosDataset(fname, "trainset", comm, var_config=None) + val_base = AdiosDataset(fname, "valset", comm, var_config=None) + test_base = AdiosDataset(fname, "testset", comm, var_config=None) + trainset = HeteroFromHomogeneousDataset(train_base, edge_dim=edge_dim) + valset = HeteroFromHomogeneousDataset(val_base, edge_dim=edge_dim) + testset = HeteroFromHomogeneousDataset(test_base, edge_dim=edge_dim) + else: + basedir = os.path.join(datadir, f"{args.modelname}.pickle") + trainset = SimplePickleDataset( + basedir=basedir, label="trainset", var_config=None + ) + valset = SimplePickleDataset(basedir=basedir, label="valset", var_config=None) + testset = SimplePickleDataset(basedir=basedir, label="testset", var_config=None) + + info( + "trainset,valset,testset size: %d %d %d" + % (len(trainset), len(valset), len(testset)) + ) + + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + + trainset = EdgeAttrDatasetAdapter(trainset, edge_dim=edge_dim) + valset = EdgeAttrDatasetAdapter(valset, edge_dim=edge_dim) + testset = EdgeAttrDatasetAdapter(testset, edge_dim=edge_dim) + + (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders( + trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"] + ) + + config = update_config(config, train_loader, val_loader, test_loader) + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + if arch_config.get("mpnn_type") == "HeteroPNA" and not arch_config.get("pna_deg"): + info("Computing pna_deg for HeteroPNA from training dataset") + pna_deg = compute_pna_deg_for_hetero_dataset(trainset, verbosity=2) + arch_config["pna_deg"] = pna_deg + arch_config["max_neighbours"] = max(0, len(pna_deg) - 1) + + config = _to_jsonable(config) + hydragnn.utils.input_config_parsing.save_config(config, log_name) + + precision = config["NeuralNetwork"]["Training"].get("precision", "fp32") + metadata = None + try: + metadata = trainset[0].metadata() + except Exception as exc: + info(f"Unable to fetch hetero metadata: {exc}") + model = hydragnn.models.create_model_config( + config=config["NeuralNetwork"], + verbosity=config["Verbosity"]["level"], + metadata=metadata, + ) + + learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"] + optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=5, min_lr=1e-5 + ) + + model, optimizer = hydragnn.utils.distributed.distributed_model_wrapper( + model, + optimizer, + config["Verbosity"]["level"], + find_unused_parameters=True, + ) + + _diag("Entering print_model") + print_model(model) + _diag("Exited print_model") + + _diag("Entering load_existing_model_config") + hydragnn.utils.model.load_existing_model_config( + model, config["NeuralNetwork"]["Training"], optimizer=optimizer + ) + _diag("Exited load_existing_model_config") + + _diag("Entering train_validate_test") + hydragnn.train.train_validate_test( + model, + optimizer, + train_loader, + val_loader, + test_loader, + writer, + scheduler, + config["NeuralNetwork"], + log_name, + config["Verbosity"]["level"], + create_plots=False, + precision=precision, + ) + _diag("Exited train_validate_test") + + hydragnn.utils.model.save_model(model, optimizer, log_name) + hydragnn.utils.profiling_and_tracing.print_timers(config["Verbosity"]["level"]) + if writer is not None: + writer.close() + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/examples/opf/train_opf_graph_output_homogeneous.py b/examples/opf/train_opf_graph_output_homogeneous.py new file mode 100644 index 000000000..8563bccb5 --- /dev/null +++ b/examples/opf/train_opf_graph_output_homogeneous.py @@ -0,0 +1,673 @@ +import os +import json +import logging +import argparse +import shutil +import subprocess +import sys +import time +from mpi4py import MPI + +import torch +import torch.distributed as dist +from torch_geometric.datasets import OPFDataset +import torch_geometric.datasets.opf as tg_opf +from __init__ import data_ops +from opf_nvme_utils import stage_case_to_nvme +from opf_solution_utils import ( + assemble_edge_attr, + resolve_edge_feature_schema, + validate_voi_node_features, + info, +) + + +def _to_jsonable(obj): + if isinstance(obj, torch.Tensor): + return obj.item() if obj.numel() == 1 else obj.tolist() + if isinstance(obj, dict): + return {k: _to_jsonable(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_to_jsonable(v) for v in obj] + return obj + + +def _patch_fast_tar_extraction(): + tar_path = shutil.which("tar") + if tar_path is None: + return + + original_extract_tar = tg_opf.extract_tar + + def _fast_extract_tar(path: str, folder: str, mode: str = "r:gz", log: bool = True): + if log: + print(f"Extracting {path}", file=sys.stderr) + try: + try: + subprocess.run( + [ + tar_path, + "--checkpoint=1000", + "--checkpoint-action=dot", + "-xzf", + path, + "-C", + folder, + ], + check=True, + ) + if log: + print("", file=sys.stderr) + return + except Exception: + subprocess.run( + [tar_path, "-xzf", path, "-C", folder], + check=True, + ) + except Exception: + original_extract_tar(path, folder, mode=mode, log=log) + + tg_opf.extract_tar = _fast_extract_tar + + +import hydragnn +from hydragnn.utils.datasets.pickledataset import ( + SimplePickleWriter, + SimplePickleDataset, +) +from hydragnn.utils.distributed import nsplit +from hydragnn.utils.model import print_model +from hydragnn.utils.input_config_parsing.config_utils import update_config + +try: + from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset +except ImportError: + AdiosWriter = None + AdiosDataset = None + + +def _prepare_sample( + data, + edge_dim=None, + edge_feature_schema=None, +): + data.y = data.objective.view(1, 1).to(torch.float32) + data.graph_attr = data.x.view(1, -1).to(torch.float32) + data, _ = assemble_edge_attr( + data, edge_dim=edge_dim, feature_schema=edge_feature_schema + ) + data_h = data.to_homogeneous( + node_attrs=["x"], + edge_attrs=["edge_attr"], + add_node_type=True, + add_edge_type=True, + ) + data_h.y = data.y + data_h.graph_attr = data.graph_attr + return data_h + + +def _load_split(root, split, case_name, num_groups, topological_perturbations): + if dist.is_initialized(): + if dist.get_rank() == 0: + info( + f"Loading OPF split: case={case_name} split={split} groups={num_groups}" + ) + else: + info(f"Loading OPF split: case={case_name} split={split} groups={num_groups}") + + def _construct(force_reload: bool = False): + return OPFDataset( + root=root, + split=split, + case_name=case_name, + num_groups=num_groups, + topological_perturbations=topological_perturbations, + force_reload=force_reload, + ) + + try: + return _construct(force_reload=False) + except Exception as exc: + msg = str(exc) + recoverable = isinstance(exc, EOFError) or any( + token in msg + for token in ( + "PytorchStreamReader failed reading file", + "PytorchStreamReader failed reading zip archive", + "failed finding central directory", + "Cannot use ``weights_only=True`` with files saved in the legacy .tar format", + "Weights only load failed", + "Unsupported operand 80", + ) + ) + if not recoverable: + raise + + if dist.is_initialized() and dist.get_world_size() > 1: + rank = dist.get_rank() + info( + f"Corrupted processed cache detected for case={case_name} split={split}; rank {rank} rebuilding with force_reload=True" + ) + dataset = _construct(force_reload=True) + MPI.COMM_WORLD.Barrier() + return dataset + + info( + f"Corrupted processed cache detected for case={case_name} split={split}; rebuilding with force_reload=True" + ) + return _construct(force_reload=True) + + +def _prime_processed_splits_on_rank0( + root, + case_name, + num_groups, + topological_perturbations, + rank, + comm, +): + if rank == 0: + for split in ("train", "val", "test"): + _load_split( + root, + split, + case_name, + num_groups, + topological_perturbations, + ) + comm.Barrier() + + +def _subset_for_rank(dataset, rank, world_size): + rx = list(nsplit(range(len(dataset)), world_size))[rank] + return [dataset[i] for i in range(rx.start, rx.stop)] + + +def _log_phase_time(comm, rank, label: str, elapsed_local: float): + elapsed_max = comm.allreduce(float(elapsed_local), op=MPI.MAX) + elapsed_sum = comm.allreduce(float(elapsed_local), op=MPI.SUM) + elapsed_avg = elapsed_sum / max(1, comm.Get_size()) + if rank == 0: + info(f"Timing {label}: max={elapsed_max:.2f}s avg={elapsed_avg:.2f}s") + + +def _resolve_preonly_case_names(args, datadir): + if not args.preonly_case_names: + return [args.case_name] + + case_names = [c.strip() for c in args.preonly_case_names if c.strip()] + if len(case_names) == 1 and case_names[0].lower() == "all": + discovered = data_ops.discover_cases(datadir, args.topological_perturbations) + if discovered: + return discovered + return [args.case_name] + return case_names + + +if __name__ == "__main__": + _patch_fast_tar_extraction() + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "--inputfile", type=str, default="opf_graph_output_homogeneous.json" + ) + parser.add_argument("--data_root", type=str, default="dataset") + parser.add_argument( + "--case_name", + type=str, + default="pglib_opf_case14_ieee", + ) + parser.add_argument( + "--num_groups", + type=str, + default="1", + help="Number of groups or 'all'", + ) + parser.add_argument( + "--num_groups_max", + type=int, + default=20, + help="Fallback/probe cap when --num_groups all and none on disk", + ) + parser.add_argument( + "--no_num_groups_probe", + action="store_false", + dest="num_groups_probe", + help="Disable probing remote storage when --num_groups all and none on disk", + ) + parser.set_defaults(num_groups_probe=True) + parser.add_argument("--topological_perturbations", action="store_true") + parser.add_argument("--preonly", action="store_true", help="preprocess only") + parser.add_argument( + "--preonly_case_names", + nargs="+", + default=None, + help="Case names to preprocess when --preonly is set; supports 'all'", + ) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--num_epoch", type=int, default=None) + parser.add_argument("--modelname", type=str, default="OPF_Homo") + parser.add_argument( + "--nvme", + action="store_true", + help="Stage selected OPF case onto node-local NVMe/scratch if available", + ) + + group = parser.add_mutually_exclusive_group() + group.add_argument("--adios", action="store_const", dest="format", const="adios") + group.add_argument("--pickle", action="store_const", dest="format", const="pickle") + parser.set_defaults(format="pickle") + + args = parser.parse_args() + + dirpwd = os.path.dirname(os.path.abspath(__file__)) + datadir = os.path.join(dirpwd, args.data_root) + input_filename = os.path.join(dirpwd, args.inputfile) + + with open(input_filename, "r") as f: + config = json.load(f) + + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + raw_edge_dim = arch_config.get("edge_dim") + if isinstance(raw_edge_dim, dict): + edge_dim = {str(k): int(v) for k, v in raw_edge_dim.items()} + edge_feature_schema = None + elif raw_edge_dim is not None: + edge_dim = int(raw_edge_dim) + names = arch_config.get("edge_feature_names") + if names: + edge_feature_schema = resolve_edge_feature_schema(names, edge_dim) + else: + edge_feature_schema = None + else: + raise RuntimeError("edge_dim must be specified in config.") + arch_config["edge_dim"] = edge_dim + + if args.batch_size is not None: + config["NeuralNetwork"]["Training"]["batch_size"] = args.batch_size + if args.num_epoch is not None: + config["NeuralNetwork"]["Training"]["num_epoch"] = args.num_epoch + + comm_size, rank = hydragnn.utils.distributed.setup_ddp() + comm = MPI.COMM_WORLD + + logging.basicConfig( + level=logging.INFO, + format=f"%(levelname)s (rank {rank}): %(message)s", + datefmt="%H:%M:%S", + ) + + log_name = args.modelname + hydragnn.utils.print.setup_log(log_name) + writer = hydragnn.utils.model.get_summary_writer(log_name) + + requested_num_groups = data_ops.parse_num_groups(args.num_groups) + + if args.preonly: + case_names = _resolve_preonly_case_names(args, datadir) + trainset, valset, testset = [], [], [] + output_datadir = datadir + + for case_name in case_names: + t_case = time.perf_counter() + num_groups = data_ops.resolve_num_groups( + requested_num_groups, + datadir, + case_name, + args.topological_perturbations, + args.num_groups_max, + args.num_groups_probe, + rank, + comm, + ) + + t_download = time.perf_counter() + data_ops.ensure_opf_downloaded( + datadir, + case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} phase=download_extract groups={num_groups}", + time.perf_counter() - t_download, + ) + + case_datadir = datadir + if args.nvme: + serialized_target = ( + f"{args.modelname}.bp" + if args.format == "adios" + else f"{args.modelname}.pickle" + ) + t_stage = time.perf_counter() + case_datadir = stage_case_to_nvme( + datadir, + case_name, + args.topological_perturbations, + comm, + rank, + None, + serialized_targets=[serialized_target], + ) + _log_phase_time( + comm, + rank, + f"case={case_name} phase=nvme_stage groups={num_groups}", + time.perf_counter() - t_stage, + ) + if args.format != "adios": + output_datadir = case_datadir + + _prime_processed_splits_on_rank0( + case_datadir, + case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + + t_load = time.perf_counter() + train_raw = _load_split( + case_datadir, + "train", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=train phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_load = time.perf_counter() + val_raw = _load_split( + case_datadir, + "val", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=val phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_load = time.perf_counter() + test_raw = _load_split( + case_datadir, + "test", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=test phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_pre = time.perf_counter() + train_subset = _subset_for_rank(train_raw, rank, comm_size) + trainset.extend( + _prepare_sample( + d, edge_dim=edge_dim, edge_feature_schema=edge_feature_schema + ) + for d in train_subset + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=train phase=preprocess local_samples={len(train_subset)}", + time.perf_counter() - t_pre, + ) + + t_pre = time.perf_counter() + val_subset = _subset_for_rank(val_raw, rank, comm_size) + valset.extend( + _prepare_sample( + d, edge_dim=edge_dim, edge_feature_schema=edge_feature_schema + ) + for d in val_subset + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=val phase=preprocess local_samples={len(val_subset)}", + time.perf_counter() - t_pre, + ) + + t_pre = time.perf_counter() + test_subset = _subset_for_rank(test_raw, rank, comm_size) + testset.extend( + _prepare_sample( + d, edge_dim=edge_dim, edge_feature_schema=edge_feature_schema + ) + for d in test_subset + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=test phase=preprocess local_samples={len(test_subset)}", + time.perf_counter() - t_pre, + ) + + _log_phase_time( + comm, + rank, + f"case={case_name} phase=prepare_total groups={num_groups}", + time.perf_counter() - t_case, + ) + + info( + f"Local split sizes: train={len(trainset)}, val={len(valset)}, test={len(testset)}" + ) + + if args.format == "adios": + t_write = time.perf_counter() + if AdiosWriter is None: + raise RuntimeError("adios2 is not available in this environment.") + fname = os.path.join(output_datadir, f"{args.modelname}.bp") + if rank == 0 and os.path.exists(fname): + if os.path.isdir(fname): + shutil.rmtree(fname, ignore_errors=True) + else: + os.remove(fname) + comm.Barrier() + adwriter = AdiosWriter(fname, comm) + adwriter.add("trainset", trainset) + adwriter.add("valset", valset) + adwriter.add("testset", testset) + adwriter.save() + _log_phase_time( + comm, + rank, + f"phase=serialize format=adios model={args.modelname}", + time.perf_counter() - t_write, + ) + else: + t_write = time.perf_counter() + basedir = os.path.join(output_datadir, f"{args.modelname}.pickle") + SimplePickleWriter(trainset, basedir, "trainset", use_subdir=True) + SimplePickleWriter(valset, basedir, "valset", use_subdir=True) + SimplePickleWriter(testset, basedir, "testset", use_subdir=True) + _log_phase_time( + comm, + rank, + f"phase=serialize format=pickle model={args.modelname}", + time.perf_counter() - t_write, + ) + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() + raise SystemExit(0) + + num_groups = data_ops.resolve_num_groups( + requested_num_groups, + datadir, + args.case_name, + args.topological_perturbations, + args.num_groups_max, + args.num_groups_probe, + rank, + comm, + ) + + data_ops.ensure_opf_downloaded( + datadir, + args.case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + + if args.nvme: + serialized_target = ( + f"{args.modelname}.bp" + if args.format == "adios" + else f"{args.modelname}.pickle" + ) + datadir = stage_case_to_nvme( + datadir, + args.case_name, + args.topological_perturbations, + comm, + rank, + None, + serialized_targets=[serialized_target], + ) + + _prime_processed_splits_on_rank0( + datadir, + args.case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + + info("Loading OPF splits...") + train_raw = _load_split( + datadir, + "train", + args.case_name, + num_groups, + args.topological_perturbations, + ) + val_raw = _load_split( + datadir, + "val", + args.case_name, + num_groups, + args.topological_perturbations, + ) + test_raw = _load_split( + datadir, + "test", + args.case_name, + num_groups, + args.topological_perturbations, + ) + + # Validate var_config from config — no auto-fill from data + var_config = config["NeuralNetwork"]["Variables_of_interest"] + validate_voi_node_features(config) + if ( + not isinstance(var_config.get("graph_feature_dims"), list) + or len(var_config["graph_feature_dims"]) == 0 + ): + raise RuntimeError( + "'graph_feature_dims' must be an explicit non-empty list in the config." + ) + + if args.format == "adios": + if AdiosDataset is None: + raise RuntimeError("adios2 is not available in this environment.") + fname = os.path.join(datadir, f"{args.modelname}.bp") + trainset = AdiosDataset(fname, "trainset", comm, var_config=var_config) + valset = AdiosDataset(fname, "valset", comm, var_config=var_config) + testset = AdiosDataset(fname, "testset", comm, var_config=var_config) + else: + basedir = os.path.join(datadir, f"{args.modelname}.pickle") + trainset = SimplePickleDataset( + basedir=basedir, label="trainset", var_config=var_config + ) + valset = SimplePickleDataset( + basedir=basedir, label="valset", var_config=var_config + ) + testset = SimplePickleDataset( + basedir=basedir, label="testset", var_config=var_config + ) + + info( + "trainset,valset,testset size: %d %d %d" + % (len(trainset), len(valset), len(testset)) + ) + + (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders( + trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"] + ) + + config = update_config(config, train_loader, val_loader, test_loader) + config = _to_jsonable(config) + hydragnn.utils.input_config_parsing.save_config(config, log_name) + + precision = config["NeuralNetwork"]["Training"].get("precision", "fp32") + model = hydragnn.models.create_model_config( + config=config["NeuralNetwork"], + verbosity=config["Verbosity"]["level"], + ) + + learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"] + optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=5, min_lr=1e-5 + ) + + model, optimizer = hydragnn.utils.distributed.distributed_model_wrapper( + model, optimizer, config["Verbosity"]["level"] + ) + + print_model(model) + + hydragnn.utils.model.load_existing_model_config( + model, config["NeuralNetwork"]["Training"], optimizer=optimizer + ) + + hydragnn.train.train_validate_test( + model, + optimizer, + train_loader, + val_loader, + test_loader, + writer, + scheduler, + config["NeuralNetwork"], + log_name, + config["Verbosity"]["level"], + create_plots=False, + precision=precision, + ) + + hydragnn.utils.model.save_model(model, optimizer, log_name) + hydragnn.utils.profiling_and_tracing.print_timers(config["Verbosity"]["level"]) + if writer is not None: + writer.close() + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/examples/opf/train_opf_solution_heterogeneous.py b/examples/opf/train_opf_solution_heterogeneous.py new file mode 100644 index 000000000..abcef6608 --- /dev/null +++ b/examples/opf/train_opf_solution_heterogeneous.py @@ -0,0 +1,1419 @@ +"""Train node-level OPF solution prediction. + +Use --max_samples to cap the total number of samples preprocessed +across train/val/test splits (proportional allocation). For example, +"--max_samples 100 --preonly" will serialize 100 total samples to pickle. + +Arguments summary: + --case_name ... Select one or more cases, or 'all'. + --num_groups Select group count or load all available groups. + --num_groups_max Fallback/probe cap when 'all' and none on disk. + --no_num_groups_probe Disable remote probing for 'all'. + --node_target_type bus|generator Choose node target type to predict. + --preonly Preprocess/serialize only (no training). + --max_samples Limit total samples across splits. + --adios / --pickle Serialization format. + --batch_size / --num_epoch Override training hyperparameters. +""" + +import os +import json +import logging +import argparse +import copy +import shutil +import subprocess +import sys +from mpi4py import MPI + +import numpy as np +import torch +import torch.distributed as dist +from torch_geometric.datasets import OPFDataset +import torch_geometric.datasets.opf as tg_opf +from __init__ import data_ops +from opf_nvme_utils import stage_case_to_nvme + + +_DEFAULT_CASE_NAMES = [ + "pglib_opf_case14_ieee", + "pglib_opf_case30_ieee", + "pglib_opf_case57_ieee", + "pglib_opf_case118_ieee", + "pglib_opf_case500_goc", + "pglib_opf_case2000_goc", + "pglib_opf_case6470_rte", + "pglib_opf_case4661_sdet", + "pglib_opf_case10000_goc", + "pglib_opf_case13659_pegase", +] + + +def _patch_fast_tar_extraction(): + tar_path = shutil.which("tar") + if tar_path is None: + return + + original_extract_tar = tg_opf.extract_tar + + def _fast_extract_tar(path: str, folder: str, mode: str = "r:gz", log: bool = True): + if log: + print(f"Extracting {path}", file=sys.stderr) + try: + try: + subprocess.run( + [ + tar_path, + "--checkpoint=1000", + "--checkpoint-action=dot", + "-xzf", + path, + "-C", + folder, + ], + check=True, + ) + if log: + print("", file=sys.stderr) + return + except Exception: + subprocess.run( + [tar_path, "-xzf", path, "-C", folder], + check=True, + ) + except Exception: + original_extract_tar(path, folder, mode=mode, log=log) + + tg_opf.extract_tar = _fast_extract_tar + + +import hydragnn +import time + + +def _diag(msg: str): + if os.getenv("HYDRAGNN_DIAG") != "1": + return + rank = dist.get_rank() if dist.is_initialized() else 0 + rank_filter = os.getenv("HYDRAGNN_DIAG_RANK") + if rank_filter is not None: + try: + if int(rank_filter) != int(rank): + return + except ValueError: + pass + now = time.perf_counter() + print(f"[diag][rank {rank}][{now:.3f}] {msg}", flush=True) + + +from hydragnn.utils.datasets.pickledataset import ( + SimplePickleWriter, + SimplePickleDataset, +) +from hydragnn.utils.distributed import nsplit +from hydragnn.utils.model import print_model +from hydragnn.utils.print import iterate_tqdm +from hydragnn.utils.input_config_parsing.config_utils import update_config + +from opf_solution_utils import ( + EdgeAttrDatasetAdapter, + OPFEnhancedModelWrapper, + HeteroFromHomogeneousDataset, + NodeBatchAdapter, + NodeTargetDatasetAdapter, + OPFDomainLoss, + assemble_edge_attr, + build_solution_target as _build_solution_target, + compute_pna_deg_for_hetero_dataset, + validate_voi_node_features, + ensure_node_y_loc as _ensure_node_y_loc, + info, + resolve_edge_feature_schema, + resolve_node_target_type as _resolve_node_target_type, +) + +from hydragnn.utils.datasets.hdf5dataset import HDF5Writer, HDF5Dataset + +try: + from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset +except ImportError: + AdiosWriter = None + AdiosDataset = None + + +def _ensure_non_scalar_attrs(data): + if hasattr(data, "_global_store"): + store = data._global_store + for key in list(store.keys()): + value = store[key] + if isinstance(value, torch.Tensor) and value.dim() == 0: + store[key] = value.reshape(1) + elif isinstance(value, np.ndarray) and value.ndim == 0: + store[key] = value.reshape(1) + return data + keys = data.keys() if callable(data.keys) else data.keys + for key in list(keys): + value = data[key] + if isinstance(value, torch.Tensor) and value.dim() == 0: + data[key] = value.reshape(1) + elif isinstance(value, np.ndarray) and value.ndim == 0: + data[key] = value.reshape(1) + return data + + +def _validate_node_stores_for_homogeneous(data): + """Crash if any node type is missing 'x' or 'y' — no silent zero-fill.""" + if not hasattr(data, "node_types"): + return + for node_type in data.node_types: + store = data[node_type] + if not hasattr(store, "x") or store.x is None: + raise RuntimeError( + f"Node type '{node_type}' is missing feature tensor 'x'. " + "All node types must have predetermined features for " + "homogeneous conversion. Refusing to auto-create zeros." + ) + if not hasattr(store, "y") or store.y is None: + raise RuntimeError( + f"Node type '{node_type}' is missing target tensor 'y'. " + "All node types must have predetermined targets for " + "homogeneous conversion. Refusing to auto-create zeros." + ) + + +def _to_jsonable(obj): + if isinstance(obj, torch.Tensor): + return obj.item() if obj.numel() == 1 else obj.tolist() + if isinstance(obj, np.ndarray): + return obj.tolist() + if isinstance(obj, np.generic): + return obj.item() + if isinstance(obj, dict): + return {k: _to_jsonable(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_to_jsonable(v) for v in obj] + return obj + + +def _raw_json_to_heterodata(filepath): + """Load a single raw OPF JSON and build a HeteroData (same as OPFDataset.process).""" + from torch_geometric.data import HeteroData + + with open(filepath) as f: + obj = json.load(f) + + grid = obj["grid"] + solution = obj["solution"] + metadata = obj["metadata"] + + data = HeteroData() + data.x = torch.tensor(grid["context"]).view(-1) + data.objective = torch.tensor(metadata["objective"]) + + data["bus"].x = torch.tensor(grid["nodes"]["bus"]) + data["bus"].y = torch.tensor(solution["nodes"]["bus"]) + data["generator"].x = torch.tensor(grid["nodes"]["generator"]) + data["generator"].y = torch.tensor(solution["nodes"]["generator"]) + data["load"].x = torch.tensor(grid["nodes"]["load"]) + data["shunt"].x = torch.tensor(grid["nodes"]["shunt"]) + + data["bus", "ac_line", "bus"].edge_index = tg_opf.extract_edge_index(obj, "ac_line") + data["bus", "ac_line", "bus"].edge_attr = torch.tensor( + grid["edges"]["ac_line"]["features"] + ) + data["bus", "ac_line", "bus"].edge_label = torch.tensor( + solution["edges"]["ac_line"]["features"] + ) + + data["bus", "transformer", "bus"].edge_index = tg_opf.extract_edge_index( + obj, "transformer" + ) + data["bus", "transformer", "bus"].edge_attr = torch.tensor( + grid["edges"]["transformer"]["features"] + ) + data["bus", "transformer", "bus"].edge_label = torch.tensor( + solution["edges"]["transformer"]["features"] + ) + + data["generator", "generator_link", "bus"].edge_index = tg_opf.extract_edge_index( + obj, "generator_link" + ) + data[ + "bus", "generator_link", "generator" + ].edge_index = tg_opf.extract_edge_index_rev(obj, "generator_link") + data["load", "load_link", "bus"].edge_index = tg_opf.extract_edge_index( + obj, "load_link" + ) + data["bus", "load_link", "load"].edge_index = tg_opf.extract_edge_index_rev( + obj, "load_link" + ) + data["shunt", "shunt_link", "bus"].edge_index = tg_opf.extract_edge_index( + obj, "shunt_link" + ) + data["bus", "shunt_link", "shunt"].edge_index = tg_opf.extract_edge_index_rev( + obj, "shunt_link" + ) + + return data + + +def _iter_raw_split_for_rank( + datadir, + case_name, + num_groups, + topological_perturbations, + split, + rank, + world_size, +): + """Yield HeteroData samples from raw JSON files for *split*, only those + assigned to *rank*. Reads one file at a time — never holds more than + one sample in memory. + """ + release = "dataset_release_1" + if topological_perturbations: + release += "_nminusone" + raw_dir = os.path.join(datadir, release, case_name, "raw") + tmp_root = os.path.join(raw_dir, "gridopt-dataset-tmp", release, case_name) + + total_samples = 15_000 * num_groups + train_limit = int(total_samples * 0.9) + val_limit = train_limit + int(total_samples * 0.05) + + # Collect all (global_index, filepath) for the requested split + split_files = [] + for group in range(num_groups): + group_dir = os.path.join(tmp_root, f"group_{group}") + for name in sorted(os.listdir(group_dir)): + i = int(name.split(".")[0].split("_")[1]) + if split == "train" and i < train_limit: + split_files.append((i, os.path.join(group_dir, name))) + elif split == "val" and train_limit <= i < val_limit: + split_files.append((i, os.path.join(group_dir, name))) + elif split == "test" and i >= val_limit: + split_files.append((i, os.path.join(group_dir, name))) + + # Sort by index for deterministic ordering + split_files.sort(key=lambda x: x[0]) + + # Select only this rank's share + n = len(split_files) + chunk = n // world_size + remainder = n % world_size + start = rank * chunk + min(rank, remainder) + end = start + chunk + (1 if rank < remainder else 0) + + skipped = 0 + for _, filepath in split_files[start:end]: + try: + yield _raw_json_to_heterodata(filepath) + except (json.JSONDecodeError, KeyError, ValueError) as exc: + skipped += 1 + logging.warning("Skipping corrupt file %s: %s", filepath, exc) + if skipped: + logging.warning("Rank skipped %d corrupt JSON file(s) in this split", skipped) + + +def _prepare_sample( + data, + node_target_type: str, + case_name: str, + to_homogeneous: bool = False, + edge_dim=None, + edge_feature_schema=None, +): + data.y = _build_solution_target(data, node_target_type) + _ensure_node_y_loc(data) + data.graph_attr = data.x.view(1, -1).to(torch.float32) + data.case_name = case_name + _ensure_non_scalar_attrs(data) + if hasattr(data, "num_nodes_dict"): + delattr(data, "num_nodes_dict") + data, _ = assemble_edge_attr( + data, edge_dim=edge_dim, feature_schema=edge_feature_schema + ) + if not to_homogeneous: + return data + _validate_node_stores_for_homogeneous(data) + data_h = data.to_homogeneous( + node_attrs=["x", "y"], + edge_attrs=["edge_attr"], + add_node_type=True, + add_edge_type=True, + ) + data_h.graph_attr = data.graph_attr + data_h.case_name = case_name + _ensure_non_scalar_attrs(data_h) + if hasattr(data_h, "num_nodes_dict"): + delattr(data_h, "num_nodes_dict") + return data_h + + +def _load_split(root, split, case_name, num_groups, topological_perturbations): + if dist.is_initialized(): + if dist.get_rank() == 0: + info( + f"Loading OPF split: case={case_name} split={split} groups={num_groups}" + ) + else: + info(f"Loading OPF split: case={case_name} split={split} groups={num_groups}") + + def _construct(force_reload: bool = False): + return OPFDataset( + root=root, + split=split, + case_name=case_name, + num_groups=num_groups, + topological_perturbations=topological_perturbations, + force_reload=force_reload, + ) + + try: + return _construct(force_reload=False) + except Exception as exc: + msg = str(exc) + recoverable = isinstance(exc, EOFError) or any( + token in msg + for token in ( + "PytorchStreamReader failed reading file", + "PytorchStreamReader failed reading zip archive", + "failed finding central directory", + "Cannot use ``weights_only=True`` with files saved in the legacy .tar format", + "Weights only load failed", + "Unsupported operand 80", + ) + ) + if not recoverable: + raise + + if dist.is_initialized() and dist.get_world_size() > 1: + rank = dist.get_rank() + info( + f"Corrupted processed cache detected for case={case_name} split={split}; rank {rank} rebuilding with force_reload=True" + ) + dataset = _construct(force_reload=True) + MPI.COMM_WORLD.Barrier() + return dataset + + info( + f"Corrupted processed cache detected for case={case_name} split={split}; rebuilding with force_reload=True" + ) + return _construct(force_reload=True) + + +def _prime_processed_splits_on_rank0( + root, + case_names, + case_num_groups, + topological_perturbations, + rank, + comm, +): + if rank == 0: + _diag("Priming OPF processed split caches on rank 0") + for case_name in case_names: + num_groups = case_num_groups[case_name] + for split in ("train", "val", "test"): + _load_split( + root, + split, + case_name, + num_groups, + topological_perturbations, + ) + _diag("Finished OPF processed split cache priming on rank 0") + comm.Barrier() + + +def _parse_case_list(case_name_args): + if not case_name_args: + return [] + if isinstance(case_name_args, str): + return [case_name_args] + return [c.strip() for c in case_name_args if c.strip()] + + +def _resolve_preonly_case_names(args, datadir, default_case_names): + requested = _parse_case_list(args.preonly_case_names) + if not requested: + return default_case_names + if len(requested) == 1 and requested[0].lower() == "all": + discovered = data_ops.discover_cases(datadir, args.topological_perturbations) + if discovered: + return discovered + return default_case_names + return requested + + +def _subset_for_rank(dataset, rank, world_size, max_samples=None): + if max_samples is None: + indices = range(len(dataset)) + else: + max_samples = max(0, min(int(max_samples), len(dataset))) + indices = range(max_samples) + rx = list(nsplit(indices, world_size))[rank] + return [dataset[i] for i in range(rx.start, rx.stop)] + + +def _allocate_split_caps(max_samples, split_sizes): + if max_samples is None: + return {name: None for name in split_sizes} + total = sum(split_sizes.values()) + if total <= 0: + return {name: 0 for name in split_sizes} + raw = { + name: (max_samples * (size / total)) if size > 0 else 0 + for name, size in split_sizes.items() + } + caps = {name: min(split_sizes[name], int(raw[name])) for name in split_sizes} + remainder = max_samples - sum(caps.values()) + order = sorted( + split_sizes.keys(), + key=lambda name: (raw[name] - int(raw[name])), + reverse=True, + ) + for name in order: + if remainder <= 0: + break + if caps[name] < split_sizes[name]: + caps[name] += 1 + remainder -= 1 + return caps + + +def _log_phase_time(comm, rank, label: str, elapsed_local: float): + elapsed_max = comm.allreduce(float(elapsed_local), op=MPI.MAX) + elapsed_sum = comm.allreduce(float(elapsed_local), op=MPI.SUM) + elapsed_avg = elapsed_sum / max(1, comm.Get_size()) + if rank == 0: + info(f"Timing {label}: max={elapsed_max:.2f}s avg={elapsed_avg:.2f}s") + + +if __name__ == "__main__": + _patch_fast_tar_extraction() + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "--inputfile", type=str, default="opf_solution_heterogeneous.json" + ) + parser.add_argument("--data_root", type=str, default="dataset") + parser.add_argument( + "--case_name", + nargs="+", + default=["pglib_opf_case14_ieee"], + help="Case name(s) or 'all'", + ) + parser.add_argument( + "--num_groups", + type=str, + default="1", + help="Number of groups or 'all'", + ) + parser.add_argument( + "--num_groups_max", + type=int, + default=20, + help="Fallback/probe cap when --num_groups all and none on disk", + ) + parser.add_argument( + "--no_num_groups_probe", + action="store_false", + dest="num_groups_probe", + help="Disable probing remote storage when --num_groups all and none on disk", + ) + parser.set_defaults(num_groups_probe=True) + parser.add_argument("--topological_perturbations", action="store_true") + parser.add_argument("--preonly", action="store_true", help="preprocess only") + parser.add_argument( + "--preonly_case_names", + nargs="+", + default=None, + help="Case names to preprocess when --preonly is set; supports 'all'", + ) + parser.add_argument( + "--max_samples", + type=int, + default=None, + help="Limit total number of samples across train/val/test splits", + ) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--num_epoch", type=int, default=None) + parser.add_argument("--modelname", type=str, default="OPF_Solution_Hetero") + parser.add_argument("--mpnn_type", type=str, default=None) + parser.add_argument("--hidden_dim", type=int, default=None) + parser.add_argument("--num_conv_layers", type=int, default=None) + parser.add_argument("--learning_rate", type=float, default=None) + parser.add_argument("--log", type=str, default=None) + domain_group = parser.add_mutually_exclusive_group() + domain_group.add_argument( + "--enable_domain_loss", + action="store_true", + help="Enable OPF domain-informed auxiliary loss regardless of config default.", + ) + domain_group.add_argument( + "--disable_domain_loss", + action="store_true", + help="Disable OPF domain-informed auxiliary loss regardless of config default.", + ) + parser.add_argument( + "--domain_loss_voltage_bound_weight", + type=float, + default=None, + help="Override DomainLoss.voltage_bound_weight.", + ) + parser.add_argument( + "--domain_loss_voltage_bound_feature_indices", + nargs=2, + type=int, + default=None, + metavar=("VMIN_IDX", "VMAX_IDX"), + help="Override DomainLoss.voltage_bound_feature_indices.", + ) + parser.add_argument( + "--domain_loss_voltage_output_index", + type=int, + default=None, + help="Override DomainLoss.voltage_output_index (index of Vm in bus_pred; default 1).", + ) + parser.add_argument( + "--domain_loss_va_output_index", + type=int, + default=None, + help="Override DomainLoss.va_output_index (index of Va in bus_pred; default 0).", + ) + parser.add_argument( + "--domain_loss_angle_diff_weight", + type=float, + default=None, + help="Override DomainLoss.angle_diff_weight (angle-difference-limit penalty).", + ) + parser.add_argument( + "--domain_loss_line_flow_weight", + type=float, + default=None, + help="Override DomainLoss.line_flow_weight (DC thermal-limit penalty).", + ) + parser.add_argument( + "--domain_loss_line_flow_slack", + type=float, + default=None, + help=( + "Override DomainLoss.line_flow_slack: tolerance subtracted from rate_a before " + "penalising, absorbing DC-approximation linearisation error (default 1e-4)." + ), + ) + parser.add_argument( + "--domain_loss_ema_momentum", + type=float, + default=None, + help="Override DomainLoss.ema_momentum for per-term EMA normalization (default 0.1).", + ) + parser.add_argument( + "--domain_loss_warmup_epochs", + type=int, + default=None, + help="Override DomainLoss.warmup_epochs: epochs with zero domain-loss weight (default 0).", + ) + parser.add_argument( + "--domain_loss_ramp_epochs", + type=int, + default=None, + help="Override DomainLoss.ramp_epochs: epochs to linearly ramp from 0 to full weight (default 0).", + ) + parser.add_argument( + "--nvme", + action="store_true", + help="Stage selected OPF case(s) onto node-local NVMe/scratch if available", + ) + parser.add_argument( + "--node_target_type", + type=str, + default="bus", + choices=["bus", "generator"], + ) + + group = parser.add_mutually_exclusive_group() + group.add_argument("--adios", action="store_const", dest="format", const="adios") + group.add_argument("--pickle", action="store_const", dest="format", const="pickle") + group.add_argument("--hdf5", action="store_const", dest="format", const="hdf5") + parser.set_defaults(format="pickle") + + args = parser.parse_args() + + dirpwd = os.path.dirname(os.path.abspath(__file__)) + datadir = os.path.join(dirpwd, args.data_root) + input_filename = os.path.join(dirpwd, args.inputfile) + + with open(input_filename, "r") as f: + config = json.load(f) + + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + + # CLI overrides for HPO + for param in ("mpnn_type", "hidden_dim", "num_conv_layers"): + val = getattr(args, param, None) + if val is not None: + arch_config[param] = val + if args.learning_rate is not None: + config["NeuralNetwork"]["Training"]["Optimizer"][ + "learning_rate" + ] = args.learning_rate + if args.batch_size is not None: + config["NeuralNetwork"]["Training"]["batch_size"] = args.batch_size + if args.num_epoch is not None: + config["NeuralNetwork"]["Training"]["num_epoch"] = args.num_epoch + + training_config = config.setdefault("NeuralNetwork", {}).setdefault("Training", {}) + + # Apply CLI overrides for domain loss. Any CLI flag takes precedence over + # whatever is stored in the input config. + _domain_cli_overrides = { + "enabled": True if args.enable_domain_loss else (False if args.disable_domain_loss else None), + "voltage_bound_weight": args.domain_loss_voltage_bound_weight, + "voltage_bound_feature_indices": ( + list(args.domain_loss_voltage_bound_feature_indices) + if args.domain_loss_voltage_bound_feature_indices is not None + else None + ), + "voltage_output_index": args.domain_loss_voltage_output_index, + "va_output_index": args.domain_loss_va_output_index, + "angle_diff_weight": args.domain_loss_angle_diff_weight, + "line_flow_weight": args.domain_loss_line_flow_weight, + "line_flow_slack": args.domain_loss_line_flow_slack, + "ema_momentum": args.domain_loss_ema_momentum, + "warmup_epochs": args.domain_loss_warmup_epochs, + "ramp_epochs": args.domain_loss_ramp_epochs, + } + if any(v is not None for v in _domain_cli_overrides.values()): + domain_loss_config = copy.deepcopy(training_config.get("DomainLoss", {})) + for key, val in _domain_cli_overrides.items(): + if val is not None: + domain_loss_config[key] = val + training_config["DomainLoss"] = domain_loss_config + + raw_edge_dim = arch_config.get("edge_dim") + if isinstance(raw_edge_dim, dict): + # Heterogeneous route: per-edge-type widths from pre-assembled tensors. + edge_dim = {str(k): int(v) for k, v in raw_edge_dim.items()} + edge_feature_schema = None + elif raw_edge_dim is not None: + # Homogeneous route: uniform width, optional named-column schema. + edge_dim = int(raw_edge_dim) + names = arch_config.get("edge_feature_names") + if names: + edge_feature_schema = resolve_edge_feature_schema(names, edge_dim) + else: + edge_feature_schema = None + else: + raise RuntimeError("edge_dim must be specified in config.") + arch_config["edge_dim"] = edge_dim + + if "node_target_type" in config.get("NeuralNetwork", {}).get("Architecture", {}): + args.node_target_type = config["NeuralNetwork"]["Architecture"][ + "node_target_type" + ] + validate_voi_node_features(config, args.node_target_type) + + comm_size, rank = hydragnn.utils.distributed.setup_ddp() + comm = MPI.COMM_WORLD + + logging.basicConfig( + level=logging.INFO, + format=f"%(levelname)s (rank {rank}): %(message)s", + datefmt="%H:%M:%S", + ) + + log_name = args.log if args.log is not None else args.modelname + hydragnn.utils.print.setup_log(log_name) + writer = hydragnn.utils.model.get_summary_writer(log_name) + + requested_num_groups = data_ops.parse_num_groups(args.num_groups) + serialized_training_only = args.format in ("adios", "hdf5") and not args.preonly + if serialized_training_only: + case_names = [] + if rank == 0: + info( + f"{args.format.upper()} training mode: skipping OPF case discovery/download/split preparation." + ) + else: + parsed_case_names = _parse_case_list(args.case_name) + if len(parsed_case_names) == 1 and parsed_case_names[0].lower() == "all": + case_names = data_ops.discover_cases( + datadir, args.topological_perturbations + ) + if not case_names: + case_names = list(_DEFAULT_CASE_NAMES) + if not case_names: + raise RuntimeError("No OPF cases found.") + else: + case_names = parsed_case_names + + if args.preonly: + case_names = _resolve_preonly_case_names(args, datadir, case_names) + + case_num_groups = {} + shared_datadir = datadir + active_datadir = datadir + _fmt_ext = {"adios": ".bp", "hdf5": ".h5", "pickle": ".pickle"} + serialized_target = f"{args.modelname}{_fmt_ext[args.format]}" + preonly_pipeline = args.preonly and args.max_samples is None + if preonly_pipeline: + store_homogeneous = args.format == "adios" + verbosity = config["Verbosity"]["level"] + # For HDF5 streaming mode we write samples as they are processed, + # so no large lists are needed. For other formats we still + # accumulate. + hdf5_streaming = args.format == "hdf5" + if hdf5_streaming: + serialize_datadir = shared_datadir + basedir = os.path.join(serialize_datadir, f"{args.modelname}.h5") + if rank == 0 and os.path.exists(basedir): + shutil.rmtree(basedir, ignore_errors=True) + comm.Barrier() + h5writer = HDF5Writer(basedir, comm) + trainset_count = 0 + valset_count = 0 + testset_count = 0 + else: + trainset = [] + valset = [] + testset = [] + + # Task-parallel: partition cases across rank groups so each group + # processes a subset of cases concurrently (embarrassingly parallel). + _task_parallel = preonly_pipeline and len(case_names) > 1 and comm_size > 1 + _task_cases = case_names + _saved_comm, _saved_rank, _saved_size = comm, rank, comm_size + _case_sub_comm = None + if _task_parallel: + _n_cases = len(case_names) + _n_groups = min(_n_cases, comm_size) + _group_id = min(rank * _n_groups // comm_size, _n_groups - 1) + _case_sub_comm = comm.Split(_group_id, rank) + _task_cases = [ + case_names[i] + for i in range(_n_cases) + if min(i * _n_groups // _n_cases, _n_groups - 1) == _group_id + ] + if rank == 0: + info( + f"Task-parallel preprocessing: {_n_cases} cases across " + f"{_n_groups} groups of {comm_size} total ranks" + ) + # Shadow comm/rank/comm_size so the loop body uses the sub-communicator + # for barriers, broadcasts, and work splitting within each case. + comm = _case_sub_comm + rank = _case_sub_comm.Get_rank() + comm_size = _case_sub_comm.Get_size() + + for case_name in _task_cases: + t_case = time.perf_counter() + num_groups = data_ops.resolve_num_groups( + requested_num_groups, + shared_datadir, + case_name, + args.topological_perturbations, + args.num_groups_max, + args.num_groups_probe, + rank, + comm, + ) + case_num_groups[case_name] = num_groups + + t_download = time.perf_counter() + data_ops.ensure_opf_downloaded( + shared_datadir, + case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} phase=download_extract groups={num_groups}", + time.perf_counter() - t_download, + ) + + if args.nvme: + t_stage = time.perf_counter() + staged_datadir = stage_case_to_nvme( + shared_datadir, + case_name, + args.topological_perturbations, + comm, + rank, + None, + serialized_targets=[serialized_target], + ) + _log_phase_time( + comm, + rank, + f"case={case_name} phase=nvme_stage groups={num_groups}", + time.perf_counter() - t_stage, + ) + if staged_datadir != shared_datadir: + active_datadir = staged_datadir + else: + active_datadir = shared_datadir + break + _log_phase_time( + comm, + rank, + f"case={case_name} phase=prepare_total groups={num_groups}", + time.perf_counter() - t_case, + ) + + if preonly_pipeline: + for split_name, label in [ + ("train", "trainset"), + ("val", "valset"), + ("test", "testset"), + ]: + t_pre = time.perf_counter() + sample_iter = _iter_raw_split_for_rank( + active_datadir, + case_name, + num_groups, + args.topological_perturbations, + split_name, + rank, + comm_size, + ) + local_count = 0 + if hdf5_streaming: + h5writer.begin(label) + for d in iterate_tqdm( + sample_iter, + verbosity, + desc=f"Preprocess {split_name} {case_name}", + leave=False, + ): + h5writer.put( + _prepare_sample( + d, + args.node_target_type, + case_name, + store_homogeneous, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + ) + local_count += 1 + h5writer.end_label() + if label == "trainset": + trainset_count += local_count + elif label == "valset": + valset_count += local_count + else: + testset_count += local_count + else: + target_list = ( + trainset + if label == "trainset" + else (valset if label == "valset" else testset) + ) + for d in iterate_tqdm( + sample_iter, + verbosity, + desc=f"Preprocess {split_name} {case_name}", + leave=False, + ): + target_list.append( + _prepare_sample( + d, + args.node_target_type, + case_name, + store_homogeneous, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + ) + local_count += 1 + _log_phase_time( + comm, + rank, + f"case={case_name} split={split_name} phase=preprocess local_samples={local_count}", + time.perf_counter() - t_pre, + ) + + datadir = active_datadir + + # Restore original communicator after task-parallel loop. + comm, rank, comm_size = _saved_comm, _saved_rank, _saved_size + if _case_sub_comm is not None: + _case_sub_comm.Free() + + if preonly_pipeline: + if hdf5_streaming: + info( + f"Local split sizes: train={trainset_count}, val={valset_count}, test={testset_count}" + ) + else: + info( + f"Local split sizes: train={len(trainset)}, val={len(valset)}, test={len(testset)}" + ) + + t_write = time.perf_counter() + if hdf5_streaming: + # Samples already written; just finalize metadata. + h5writer.save() + elif args.format == "adios": + if AdiosWriter is None: + raise RuntimeError("adios2 is not available in this environment.") + serialize_datadir = shared_datadir + fname = os.path.join(serialize_datadir, f"{args.modelname}.bp") + if rank == 0 and os.path.exists(fname): + if os.path.isdir(fname): + shutil.rmtree(fname, ignore_errors=True) + else: + os.remove(fname) + comm.Barrier() + adwriter = AdiosWriter(fname, comm) + adwriter.add("trainset", trainset) + adwriter.add("valset", valset) + adwriter.add("testset", testset) + adwriter.save() + elif args.format == "hdf5": + serialize_datadir = shared_datadir + basedir = os.path.join(serialize_datadir, f"{args.modelname}.h5") + if rank == 0 and os.path.exists(basedir): + shutil.rmtree(basedir, ignore_errors=True) + comm.Barrier() + h5writer = HDF5Writer(basedir, comm) + h5writer.add("trainset", trainset) + h5writer.add("valset", valset) + h5writer.add("testset", testset) + h5writer.save() + else: + basedir = os.path.join(datadir, f"{args.modelname}.pickle") + SimplePickleWriter(trainset, basedir, "trainset", use_subdir=True) + SimplePickleWriter(valset, basedir, "valset", use_subdir=True) + SimplePickleWriter(testset, basedir, "testset", use_subdir=True) + _log_phase_time( + comm, + rank, + f"phase=serialize format={args.format} model={args.modelname}", + time.perf_counter() - t_write, + ) + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() + raise SystemExit(0) + + if case_names: + _prime_processed_splits_on_rank0( + datadir, + case_names, + case_num_groups, + args.topological_perturbations, + rank, + comm, + ) + + if rank == 0 and case_names: + info("Loading OPF splits...") + train_raw = [] + val_raw = [] + test_raw = [] + for case_name in case_names: + num_groups = case_num_groups[case_name] + + t_load = time.perf_counter() + train_raw.append( + _load_split( + datadir, + "train", + case_name, + num_groups, + args.topological_perturbations, + ) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=train phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_load = time.perf_counter() + val_raw.append( + _load_split( + datadir, + "val", + case_name, + num_groups, + args.topological_perturbations, + ) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=val phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_load = time.perf_counter() + test_raw.append( + _load_split( + datadir, + "test", + case_name, + num_groups, + args.topological_perturbations, + ) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=test phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + split_sizes = { + "train": sum(len(d) for d in train_raw), + "val": sum(len(d) for d in val_raw), + "test": sum(len(d) for d in test_raw), + } + split_caps = _allocate_split_caps(args.max_samples, split_sizes) + if args.max_samples is not None: + info( + "Limiting samples across splits: " + f"train={split_caps['train']}, val={split_caps['val']}, test={split_caps['test']}" + ) + + if args.preonly: + store_homogeneous = args.format == "adios" + verbosity = config["Verbosity"]["level"] + trainset = [] + valset = [] + testset = [] + remaining_caps = dict(split_caps) + for case_name, train_split in zip(case_names, train_raw): + t_pre = time.perf_counter() + case_cap = remaining_caps["train"] + if case_cap is not None: + case_cap = max(0, min(case_cap, len(train_split))) + subset = _subset_for_rank(train_split, rank, comm_size, case_cap) + for d in iterate_tqdm( + subset, verbosity, desc=f"Preprocess train {case_name}", leave=False + ): + trainset.append( + _prepare_sample( + d, + args.node_target_type, + case_name, + store_homogeneous, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + ) + if remaining_caps["train"] is not None: + remaining_caps["train"] = max( + 0, remaining_caps["train"] - min(len(train_split), case_cap or 0) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=train phase=preprocess local_samples={len(subset)}", + time.perf_counter() - t_pre, + ) + for case_name, val_split in zip(case_names, val_raw): + t_pre = time.perf_counter() + case_cap = remaining_caps["val"] + if case_cap is not None: + case_cap = max(0, min(case_cap, len(val_split))) + subset = _subset_for_rank(val_split, rank, comm_size, case_cap) + for d in iterate_tqdm( + subset, verbosity, desc=f"Preprocess val {case_name}", leave=False + ): + valset.append( + _prepare_sample( + d, + args.node_target_type, + case_name, + store_homogeneous, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + ) + if remaining_caps["val"] is not None: + remaining_caps["val"] = max( + 0, remaining_caps["val"] - min(len(val_split), case_cap or 0) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=val phase=preprocess local_samples={len(subset)}", + time.perf_counter() - t_pre, + ) + for case_name, test_split in zip(case_names, test_raw): + t_pre = time.perf_counter() + case_cap = remaining_caps["test"] + if case_cap is not None: + case_cap = max(0, min(case_cap, len(test_split))) + subset = _subset_for_rank(test_split, rank, comm_size, case_cap) + for d in iterate_tqdm( + subset, verbosity, desc=f"Preprocess test {case_name}", leave=False + ): + testset.append( + _prepare_sample( + d, + args.node_target_type, + case_name, + store_homogeneous, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + ) + if remaining_caps["test"] is not None: + remaining_caps["test"] = max( + 0, remaining_caps["test"] - min(len(test_split), case_cap or 0) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=test phase=preprocess local_samples={len(subset)}", + time.perf_counter() - t_pre, + ) + + info( + f"Local split sizes: train={len(trainset)}, val={len(valset)}, test={len(testset)}" + ) + + t_write = time.perf_counter() + if args.format == "adios": + if AdiosWriter is None: + raise RuntimeError("adios2 is not available in this environment.") + serialize_datadir = shared_datadir + fname = os.path.join(serialize_datadir, f"{args.modelname}.bp") + if rank == 0 and os.path.exists(fname): + if os.path.isdir(fname): + shutil.rmtree(fname, ignore_errors=True) + else: + os.remove(fname) + comm.Barrier() + adwriter = AdiosWriter(fname, comm) + adwriter.add("trainset", trainset) + adwriter.add("valset", valset) + adwriter.add("testset", testset) + adwriter.save() + elif args.format == "hdf5": + serialize_datadir = shared_datadir + basedir = os.path.join(serialize_datadir, f"{args.modelname}.h5") + if rank == 0 and os.path.exists(basedir): + shutil.rmtree(basedir, ignore_errors=True) + comm.Barrier() + h5writer = HDF5Writer(basedir, comm) + h5writer.add("trainset", trainset) + h5writer.add("valset", valset) + h5writer.add("testset", testset) + h5writer.save() + else: + basedir = os.path.join(datadir, f"{args.modelname}.pickle") + SimplePickleWriter(trainset, basedir, "trainset", use_subdir=True) + SimplePickleWriter(valset, basedir, "valset", use_subdir=True) + SimplePickleWriter(testset, basedir, "testset", use_subdir=True) + _log_phase_time( + comm, + rank, + f"phase=serialize format={args.format} model={args.modelname}", + time.perf_counter() - t_write, + ) + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() + raise SystemExit(0) + + if args.format == "adios": + if AdiosDataset is None: + raise RuntimeError("adios2 is not available in this environment.") + fname = os.path.join(datadir, f"{args.modelname}.bp") + if serialized_training_only and not os.path.isdir(fname): + raise RuntimeError( + f"Expected preprocessed ADIOS dataset at '{fname}' for training-only mode. " + "Run with --preonly --adios first." + ) + train_base = AdiosDataset(fname, "trainset", comm, var_config=None) + val_base = AdiosDataset(fname, "valset", comm, var_config=None) + test_base = AdiosDataset(fname, "testset", comm, var_config=None) + trainset = HeteroFromHomogeneousDataset(train_base, edge_dim=edge_dim) + valset = HeteroFromHomogeneousDataset(val_base, edge_dim=edge_dim) + testset = HeteroFromHomogeneousDataset(test_base, edge_dim=edge_dim) + elif args.format == "hdf5": + basedir = os.path.join(datadir, f"{args.modelname}.h5") + if serialized_training_only and not os.path.isdir(basedir): + raise RuntimeError( + f"Expected preprocessed HDF5 dataset at '{basedir}' for training-only mode. " + "Run with --preonly --hdf5 first." + ) + trainset = HDF5Dataset(basedir, "trainset") + valset = HDF5Dataset(basedir, "valset") + testset = HDF5Dataset(basedir, "testset") + else: + basedir = os.path.join(datadir, f"{args.modelname}.pickle") + trainset = SimplePickleDataset( + basedir=basedir, label="trainset", var_config=None + ) + valset = SimplePickleDataset(basedir=basedir, label="valset", var_config=None) + testset = SimplePickleDataset(basedir=basedir, label="testset", var_config=None) + + resolved_node_target_type = _resolve_node_target_type( + trainset[0], args.node_target_type + ) + if resolved_node_target_type != args.node_target_type: + info( + f"Resolved node_target_type '{args.node_target_type}' -> '{resolved_node_target_type}'" + ) + args.node_target_type = resolved_node_target_type + config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {})[ + "node_target_type" + ] = args.node_target_type + validate_voi_node_features(config, args.node_target_type) + + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + + trainset = EdgeAttrDatasetAdapter(trainset, edge_dim=edge_dim) + valset = EdgeAttrDatasetAdapter(valset, edge_dim=edge_dim) + testset = EdgeAttrDatasetAdapter(testset, edge_dim=edge_dim) + + trainset = NodeTargetDatasetAdapter( + trainset, args.node_target_type, edge_dim=edge_dim + ) + valset = NodeTargetDatasetAdapter(valset, args.node_target_type, edge_dim=edge_dim) + testset = NodeTargetDatasetAdapter( + testset, args.node_target_type, edge_dim=edge_dim + ) + + info( + "trainset,valset,testset size: %d %d %d" + % (len(trainset), len(valset), len(testset)) + ) + + (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders( + trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"] + ) + + train_loader = NodeBatchAdapter( + train_loader, args.node_target_type, edge_dim=edge_dim + ) + val_loader = NodeBatchAdapter(val_loader, args.node_target_type, edge_dim=edge_dim) + test_loader = NodeBatchAdapter( + test_loader, args.node_target_type, edge_dim=edge_dim + ) + + config = update_config(config, train_loader, val_loader, test_loader) + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + if arch_config.get("mpnn_type") == "HeteroPNA" and not arch_config.get("pna_deg"): + info("Computing pna_deg for HeteroPNA from training dataset") + pna_deg = compute_pna_deg_for_hetero_dataset(trainset, verbosity=2) + arch_config["pna_deg"] = pna_deg + arch_config["max_neighbours"] = max(0, len(pna_deg) - 1) + + config = _to_jsonable(config) + hydragnn.utils.input_config_parsing.save_config(config, log_name) + + precision = config["NeuralNetwork"]["Training"].get("precision", "fp32") + metadata = None + try: + metadata = trainset[0].metadata() + except Exception as exc: + info(f"Unable to fetch hetero metadata: {exc}") + node_input_dims = ( + config.get("NeuralNetwork", {}).get("Architecture", {}).get("node_input_dims") + ) + if node_input_dims is None: + raise RuntimeError( + "Missing NeuralNetwork.Architecture.node_input_dims in config. " + "Add node_input_dims to the config to initialize node embedders." + ) + model = hydragnn.models.create_model_config( + config=config["NeuralNetwork"], + verbosity=config["Verbosity"]["level"], + metadata=metadata, + node_input_dims=node_input_dims, + ) + + domain_loss_config = config["NeuralNetwork"]["Training"].get("DomainLoss") + if domain_loss_config is not None: + dl_enabled = domain_loss_config.get("enabled", False) + if rank == 0: + info( + f"[DomainLoss] config (enabled={dl_enabled}): " + + ", ".join( + f"{k}={v}" + for k, v in domain_loss_config.items() + if k != "enabled" + ) + ) + if dl_enabled and rank == 0: + info("[DomainLoss] Wrapping model with OPFEnhancedModelWrapper.") + model = OPFEnhancedModelWrapper( + model, + OPFDomainLoss( + domain_loss_config, + node_target_type=args.node_target_type, + ), + ) + + learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"] + optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=5, min_lr=1e-5 + ) + + model, optimizer = hydragnn.utils.distributed.distributed_model_wrapper( + model, + optimizer, + config["Verbosity"]["level"], + find_unused_parameters=True, + ) + + _diag("Entering print_model") + print_model(model) + _diag("Exited print_model") + + _diag("Entering load_existing_model_config") + hydragnn.utils.model.load_existing_model_config( + model, config["NeuralNetwork"]["Training"], optimizer=optimizer + ) + _diag("Exited load_existing_model_config") + + _diag("Entering train_validate_test") + hydragnn.train.train_validate_test( + model, + optimizer, + train_loader, + val_loader, + test_loader, + writer, + scheduler, + config["NeuralNetwork"], + log_name, + config["Verbosity"]["level"], + create_plots=False, + precision=precision, + ) + _diag("Exited train_validate_test") + + # Flush the final epoch's LossBreakdown line. The wrapper only flushes on + # epoch *transitions* detected inside loss(), so the last epoch's stats would + # otherwise never be written (no subsequent epoch triggers the flush). + if isinstance(model, OPFEnhancedModelWrapper): + model._flush_epoch_log(model._last_seen_epoch) + elif hasattr(model, "module") and isinstance(model.module, OPFEnhancedModelWrapper): + # DDP wraps the model in model.module + model.module._flush_epoch_log(model.module._last_seen_epoch) + + hydragnn.utils.model.save_model(model, optimizer, log_name) + hydragnn.utils.profiling_and_tracing.print_timers(config["Verbosity"]["level"]) + if writer is not None: + writer.close() + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/examples/opf/train_opf_solution_homogeneous.py b/examples/opf/train_opf_solution_homogeneous.py new file mode 100644 index 000000000..04ed955ce --- /dev/null +++ b/examples/opf/train_opf_solution_homogeneous.py @@ -0,0 +1,891 @@ +"""Train node-level OPF solution prediction (homogeneous graph). + +Arguments summary: + --case_name Select a single case or load all cases. + --case_names Comma-separated case list (used if --case_name all). + --case_list_file File with one case name per line. + --num_groups Select group count or load all available groups. + --num_groups_max Fallback/probe cap when 'all' and none on disk. + --no_num_groups_probe Disable remote probing for 'all'. + --node_target_type bus|generator Choose node target type to predict. + --preonly Preprocess/serialize only (no training). + --adios / --pickle Serialization format. + --batch_size / --num_epoch Override training hyperparameters. +""" + +import os +import json +import logging +import argparse +import shutil +import subprocess +import sys +from mpi4py import MPI + +import torch +import torch.distributed as dist +from torch_geometric.datasets import OPFDataset +import torch_geometric.datasets.opf as tg_opf +from __init__ import data_ops +from opf_nvme_utils import stage_case_to_nvme +from opf_solution_utils import ( + assemble_edge_attr, + build_solution_target as _build_solution_target, + ensure_node_y_loc as _ensure_node_y_loc, + info, + resolve_edge_feature_schema, + resolve_node_target_type as _resolve_node_target_type, +) + + +def _to_jsonable(obj): + if isinstance(obj, torch.Tensor): + return obj.item() if obj.numel() == 1 else obj.tolist() + if isinstance(obj, dict): + return {k: _to_jsonable(v) for k, v in obj.items()} + if isinstance(obj, (list, tuple)): + return [_to_jsonable(v) for v in obj] + return obj + + +_DEFAULT_CASE_NAMES = [ + "pglib_opf_case14_ieee", + "pglib_opf_case30_ieee", + "pglib_opf_case57_ieee", + "pglib_opf_case118_ieee", + "pglib_opf_case500_goc", + "pglib_opf_case2000_goc", + "pglib_opf_case6470_rte", + "pglib_opf_case4661_sdet", + "pglib_opf_case10000_goc", + "pglib_opf_case13659_pegase", +] + + +def _patch_fast_tar_extraction(): + tar_path = shutil.which("tar") + if tar_path is None: + return + + original_extract_tar = tg_opf.extract_tar + + def _fast_extract_tar(path: str, folder: str, mode: str = "r:gz", log: bool = True): + if log: + print(f"Extracting {path}", file=sys.stderr) + try: + try: + subprocess.run( + [ + tar_path, + "--checkpoint=1000", + "--checkpoint-action=dot", + "-xzf", + path, + "-C", + folder, + ], + check=True, + ) + if log: + print("", file=sys.stderr) + return + except Exception: + subprocess.run([tar_path, "-xzf", path, "-C", folder], check=True) + except Exception: + original_extract_tar(path, folder, mode=mode, log=log) + + tg_opf.extract_tar = _fast_extract_tar + + +import hydragnn +import time + + +def _diag(msg: str): + if os.getenv("HYDRAGNN_DIAG") != "1": + return + rank = dist.get_rank() if dist.is_initialized() else 0 + rank_filter = os.getenv("HYDRAGNN_DIAG_RANK") + if rank_filter is not None: + try: + if int(rank_filter) != int(rank): + return + except ValueError: + pass + now = time.perf_counter() + print(f"[diag][rank {rank}][{now:.3f}] {msg}", flush=True) + + +from hydragnn.utils.datasets.pickledataset import ( + SimplePickleWriter, + SimplePickleDataset, +) +from hydragnn.utils.distributed import nsplit +from hydragnn.utils.model import print_model +from hydragnn.utils.print import iterate_tqdm +from hydragnn.utils.input_config_parsing.config_utils import update_config + +try: + from hydragnn.utils.datasets.adiosdataset import AdiosWriter, AdiosDataset +except ImportError: + AdiosWriter = None + AdiosDataset = None + + +def _prepare_sample( + data, + node_target_type: str, + case_name: str, + to_homogeneous: bool = True, + edge_dim=None, + edge_feature_schema=None, +): + data.y = _build_solution_target(data, node_target_type) + _ensure_node_y_loc(data) + data.graph_attr = data.x.view(1, -1).to(torch.float32) + data.case_name = case_name + data, _ = assemble_edge_attr( + data, edge_dim=edge_dim, feature_schema=edge_feature_schema + ) + if not to_homogeneous: + return data + data_h = data.to_homogeneous( + node_attrs=["x", "y"], + edge_attrs=["edge_attr"], + add_node_type=True, + add_edge_type=True, + ) + # to_homogeneous fills y with NaN for node types without targets (load, shunt). + # Replace NaN with 0 so the loss doesn't blow up. + if hasattr(data_h, "y") and data_h.y is not None: + data_h.y = torch.nan_to_num(data_h.y, nan=0.0) + data_h.graph_attr = data.graph_attr + data_h.case_name = case_name + return data_h + + +def _load_split(root, split, case_name, num_groups, topological_perturbations): + if dist.is_initialized(): + if dist.get_rank() == 0: + info( + f"Loading OPF split: case={case_name} split={split} groups={num_groups}" + ) + else: + info(f"Loading OPF split: case={case_name} split={split} groups={num_groups}") + + def _construct(force_reload: bool = False): + return OPFDataset( + root=root, + split=split, + case_name=case_name, + num_groups=num_groups, + topological_perturbations=topological_perturbations, + force_reload=force_reload, + ) + + try: + return _construct(force_reload=False) + except Exception as exc: + msg = str(exc) + recoverable = isinstance(exc, EOFError) or any( + token in msg + for token in ( + "PytorchStreamReader failed reading file", + "PytorchStreamReader failed reading zip archive", + "failed finding central directory", + "Cannot use ``weights_only=True`` with files saved in the legacy .tar format", + "Weights only load failed", + "Unsupported operand 80", + ) + ) + if not recoverable: + raise + + if dist.is_initialized() and dist.get_world_size() > 1: + rank = dist.get_rank() + info( + f"Corrupted processed cache detected for case={case_name} split={split}; rank {rank} rebuilding with force_reload=True" + ) + dataset = _construct(force_reload=True) + MPI.COMM_WORLD.Barrier() + return dataset + + info( + f"Corrupted processed cache detected for case={case_name} split={split}; rebuilding with force_reload=True" + ) + return _construct(force_reload=True) + + +def _prime_processed_splits_on_rank0( + root, + case_names, + case_num_groups, + topological_perturbations, + rank, + comm, +): + if rank == 0: + for case_name in case_names: + num_groups = case_num_groups[case_name] + for split in ("train", "val", "test"): + _load_split( + root, + split, + case_name, + num_groups, + topological_perturbations, + ) + comm.Barrier() + + +def _parse_case_list(args): + cases = [] + if args.case_names: + cases.extend([c.strip() for c in args.case_names.split(",") if c.strip()]) + if args.case_list_file: + with open(args.case_list_file, "r") as f: + for line in f: + name = line.strip() + if name: + cases.append(name) + return cases + + +def _parse_preonly_case_list(args): + cases = [] + if args.preonly_case_names: + cases.extend( + [c.strip() for c in args.preonly_case_names.split(",") if c.strip()] + ) + if args.preonly_case_list_file: + with open(args.preonly_case_list_file, "r") as f: + for line in f: + name = line.strip() + if name: + cases.append(name) + return cases + + +def _subset_for_rank(dataset, rank, world_size): + rx = list(nsplit(range(len(dataset)), world_size))[rank] + return [dataset[i] for i in range(rx.start, rx.stop)] + + +def _log_phase_time(comm, rank, label: str, elapsed_local: float): + elapsed_max = comm.allreduce(float(elapsed_local), op=MPI.MAX) + elapsed_sum = comm.allreduce(float(elapsed_local), op=MPI.SUM) + elapsed_avg = elapsed_sum / max(1, comm.Get_size()) + if rank == 0: + info(f"Timing {label}: max={elapsed_max:.2f}s avg={elapsed_avg:.2f}s") + + +class HomogeneousDatasetAdapter: + def __init__(self, base, node_target_type: str): + self.base = base + self.node_target_type = node_target_type + + def __len__(self): + return len(self.base) + + def __getitem__(self, idx): + data = self.base[idx] + if hasattr(data, "node_types"): + data = data.to_homogeneous( + node_attrs=["x", "y"], + edge_attrs=["edge_attr"], + add_node_type=True, + add_edge_type=True, + ) + if not hasattr(data, "node_type") or not hasattr(data, "_node_type_names"): + raise RuntimeError("Expected homogeneous OPF sample with node_type.") + if self.node_target_type not in data._node_type_names: + raise RuntimeError( + f"Node type '{self.node_target_type}' not found in OPF sample." + ) + if not hasattr(data, "y") or data.y is None: + raise RuntimeError( + f"No targets found for node type '{self.node_target_type}' in OPF sample." + ) + # Keep y for ALL nodes — the model predicts for all nodes in the + # homogeneous graph. Non-target node types (load, shunt) have NaN + # from to_homogeneous; replace with 0 so the loss stays finite. + data.y = torch.nan_to_num(data.y, nan=0.0) + _ensure_node_y_loc(data) + return data + + def __getattr__(self, name): + return getattr(self.base, name) + + +class HomogeneousBatchAdapter: + def __init__(self, loader): + self.loader = loader + self.dataset = loader.dataset + self.sampler = getattr(loader, "sampler", None) + + def __iter__(self): + for data in self.loader: + _ensure_node_y_loc(data) + yield data + + def __len__(self): + return len(self.loader) + + def __getattr__(self, name): + return getattr(self.loader, name) + + +if __name__ == "__main__": + _patch_fast_tar_extraction() + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter + ) + parser.add_argument( + "--inputfile", type=str, default="opf_solution_homogeneous.json" + ) + parser.add_argument("--data_root", type=str, default="dataset") + parser.add_argument( + "--case_name", + type=str, + default="pglib_opf_case14_ieee", + help="Case name or 'all'", + ) + parser.add_argument( + "--case_names", + type=str, + default="", + help="Comma-separated case list (used if --case_name all)", + ) + parser.add_argument( + "--case_list_file", + type=str, + default="", + help="File with one case name per line (used if --case_name all)", + ) + parser.add_argument( + "--num_groups", + type=str, + default="1", + help="Number of groups or 'all'", + ) + parser.add_argument( + "--num_groups_max", + type=int, + default=20, + help="Fallback/probe cap when --num_groups all and none on disk", + ) + parser.add_argument( + "--no_num_groups_probe", + action="store_false", + dest="num_groups_probe", + help="Disable probing remote storage when --num_groups all and none on disk", + ) + parser.set_defaults(num_groups_probe=True) + parser.add_argument("--topological_perturbations", action="store_true") + parser.add_argument("--preonly", action="store_true", help="preprocess only") + parser.add_argument( + "--preonly_case_names", + type=str, + default="", + help="Comma-separated case list used only with --preonly", + ) + parser.add_argument( + "--preonly_case_list_file", + type=str, + default="", + help="File with one case name per line used only with --preonly", + ) + parser.add_argument("--batch_size", type=int, default=None) + parser.add_argument("--num_epoch", type=int, default=None) + parser.add_argument("--modelname", type=str, default="OPF_Solution_Homo") + parser.add_argument( + "--nvme", + action="store_true", + help="Stage selected OPF case(s) onto node-local NVMe/scratch if available", + ) + parser.add_argument( + "--node_target_type", + type=str, + default="bus", + choices=["bus", "generator"], + ) + + group = parser.add_mutually_exclusive_group() + group.add_argument("--adios", action="store_const", dest="format", const="adios") + group.add_argument("--pickle", action="store_const", dest="format", const="pickle") + parser.set_defaults(format="pickle") + + args = parser.parse_args() + + dirpwd = os.path.dirname(os.path.abspath(__file__)) + datadir = os.path.join(dirpwd, args.data_root) + input_filename = os.path.join(dirpwd, args.inputfile) + + with open(input_filename, "r") as f: + config = json.load(f) + + arch_config = config.setdefault("NeuralNetwork", {}).setdefault("Architecture", {}) + raw_edge_dim = arch_config.get("edge_dim") + if isinstance(raw_edge_dim, dict): + edge_dim = {str(k): int(v) for k, v in raw_edge_dim.items()} + edge_feature_schema = None + elif raw_edge_dim is not None: + edge_dim = int(raw_edge_dim) + names = arch_config.get("edge_feature_names") + if names: + edge_feature_schema = resolve_edge_feature_schema(names, edge_dim) + else: + edge_feature_schema = None + else: + raise RuntimeError("edge_dim must be specified in config.") + arch_config["edge_dim"] = edge_dim + + if "node_target_type" in config.get("NeuralNetwork", {}).get("Architecture", {}): + args.node_target_type = config["NeuralNetwork"]["Architecture"][ + "node_target_type" + ] + + comm_size, rank = hydragnn.utils.distributed.setup_ddp() + comm = MPI.COMM_WORLD + + logging.basicConfig( + level=logging.INFO, + format=f"%(levelname)s (rank {rank}): %(message)s", + datefmt="%H:%M:%S", + ) + + log_name = args.modelname + hydragnn.utils.print.setup_log(log_name) + writer = hydragnn.utils.model.get_summary_writer(log_name) + + requested_num_groups = data_ops.parse_num_groups(args.num_groups) + adios_training_only = args.format == "adios" and not args.preonly + if adios_training_only: + case_names = [] + if rank == 0: + info( + "ADIOS training mode: skipping OPF case discovery/download/split preparation." + ) + else: + if args.case_name.lower() == "all": + case_names = _parse_case_list(args) + if not case_names: + case_names = list(_DEFAULT_CASE_NAMES) + if not case_names: + raise RuntimeError( + "No OPF cases found. Provide --case_names or --case_list_file." + ) + else: + case_names = [args.case_name] + + if args.preonly: + preonly_case_names = _parse_preonly_case_list(args) + if len(preonly_case_names) == 1 and preonly_case_names[0].lower() == "all": + discovered = data_ops.discover_cases( + datadir, args.topological_perturbations + ) + if discovered: + preonly_case_names = discovered + else: + preonly_case_names = [] + if preonly_case_names: + case_names = preonly_case_names + + case_num_groups = {} + shared_datadir = datadir + active_datadir = datadir + serialized_target = ( + f"{args.modelname}.bp" if args.format == "adios" else f"{args.modelname}.pickle" + ) + preonly_pipeline = args.preonly + if preonly_pipeline: + verbosity = config["Verbosity"]["level"] + trainset = [] + valset = [] + testset = [] + + for case_name in case_names: + t_case = time.perf_counter() + num_groups = data_ops.resolve_num_groups( + requested_num_groups, + shared_datadir, + case_name, + args.topological_perturbations, + args.num_groups_max, + args.num_groups_probe, + rank, + comm, + ) + case_num_groups[case_name] = num_groups + + t_download = time.perf_counter() + data_ops.ensure_opf_downloaded( + shared_datadir, + case_name, + num_groups, + args.topological_perturbations, + rank, + comm, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} phase=download_extract groups={num_groups}", + time.perf_counter() - t_download, + ) + + if args.nvme: + t_stage = time.perf_counter() + staged_datadir = stage_case_to_nvme( + shared_datadir, + case_name, + args.topological_perturbations, + comm, + rank, + None, + serialized_targets=[serialized_target], + ) + _log_phase_time( + comm, + rank, + f"case={case_name} phase=nvme_stage groups={num_groups}", + time.perf_counter() - t_stage, + ) + if staged_datadir != shared_datadir: + active_datadir = staged_datadir + else: + active_datadir = shared_datadir + break + + _log_phase_time( + comm, + rank, + f"case={case_name} phase=prepare_total groups={num_groups}", + time.perf_counter() - t_case, + ) + + if preonly_pipeline: + t_load = time.perf_counter() + train_split = _load_split( + active_datadir, + "train", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=train phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_pre = time.perf_counter() + subset = _subset_for_rank(train_split, rank, comm_size) + for d in iterate_tqdm( + subset, verbosity, desc=f"Preprocess train {case_name}", leave=False + ): + trainset.append( + _prepare_sample( + d, + args.node_target_type, + case_name, + True, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=train phase=preprocess local_samples={len(subset)}", + time.perf_counter() - t_pre, + ) + + t_load = time.perf_counter() + val_split = _load_split( + active_datadir, + "val", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=val phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_pre = time.perf_counter() + subset = _subset_for_rank(val_split, rank, comm_size) + for d in iterate_tqdm( + subset, verbosity, desc=f"Preprocess val {case_name}", leave=False + ): + valset.append( + _prepare_sample( + d, + args.node_target_type, + case_name, + True, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=val phase=preprocess local_samples={len(subset)}", + time.perf_counter() - t_pre, + ) + + t_load = time.perf_counter() + test_split = _load_split( + active_datadir, + "test", + case_name, + num_groups, + args.topological_perturbations, + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=test phase=load groups={num_groups}", + time.perf_counter() - t_load, + ) + + t_pre = time.perf_counter() + subset = _subset_for_rank(test_split, rank, comm_size) + for d in iterate_tqdm( + subset, verbosity, desc=f"Preprocess test {case_name}", leave=False + ): + testset.append( + _prepare_sample( + d, + args.node_target_type, + case_name, + True, + edge_dim=edge_dim, + edge_feature_schema=edge_feature_schema, + ) + ) + _log_phase_time( + comm, + rank, + f"case={case_name} split=test phase=preprocess local_samples={len(subset)}", + time.perf_counter() - t_pre, + ) + + datadir = active_datadir + + if preonly_pipeline: + info( + f"Local split sizes: train={len(trainset)}, val={len(valset)}, test={len(testset)}" + ) + + if args.format == "adios": + t_write = time.perf_counter() + if AdiosWriter is None: + raise RuntimeError("adios2 is not available in this environment.") + serialize_datadir = shared_datadir + fname = os.path.join(serialize_datadir, f"{args.modelname}.bp") + if rank == 0 and os.path.exists(fname): + if os.path.isdir(fname): + shutil.rmtree(fname, ignore_errors=True) + else: + os.remove(fname) + comm.Barrier() + adwriter = AdiosWriter(fname, comm) + adwriter.add("trainset", trainset) + adwriter.add("valset", valset) + adwriter.add("testset", testset) + adwriter.save() + _log_phase_time( + comm, + rank, + f"phase=serialize format=adios model={args.modelname}", + time.perf_counter() - t_write, + ) + else: + t_write = time.perf_counter() + basedir = os.path.join(datadir, f"{args.modelname}.pickle") + SimplePickleWriter(trainset, basedir, "trainset", use_subdir=True) + SimplePickleWriter(valset, basedir, "valset", use_subdir=True) + SimplePickleWriter(testset, basedir, "testset", use_subdir=True) + _log_phase_time( + comm, + rank, + f"phase=serialize format=pickle model={args.modelname}", + time.perf_counter() - t_write, + ) + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() + raise SystemExit(0) + + if case_names: + _prime_processed_splits_on_rank0( + datadir, + case_names, + case_num_groups, + args.topological_perturbations, + rank, + comm, + ) + + if rank == 0 and case_names: + info("Loading OPF splits...") + train_raw = [] + val_raw = [] + test_raw = [] + for case_name in case_names: + num_groups = case_num_groups[case_name] + train_raw.append( + _load_split( + datadir, + "train", + case_name, + num_groups, + args.topological_perturbations, + ) + ) + val_raw.append( + _load_split( + datadir, + "val", + case_name, + num_groups, + args.topological_perturbations, + ) + ) + test_raw.append( + _load_split( + datadir, + "test", + case_name, + num_groups, + args.topological_perturbations, + ) + ) + + if args.format == "adios": + if AdiosDataset is None: + raise RuntimeError("adios2 is not available in this environment.") + fname = os.path.join(datadir, f"{args.modelname}.bp") + if adios_training_only and not os.path.isdir(fname): + raise RuntimeError( + f"Expected preprocessed ADIOS dataset at '{fname}' for training-only mode. " + "Run with --preonly --adios first." + ) + train_base = AdiosDataset(fname, "trainset", comm, var_config=None) + val_base = AdiosDataset(fname, "valset", comm, var_config=None) + test_base = AdiosDataset(fname, "testset", comm, var_config=None) + trainset = HomogeneousDatasetAdapter(train_base, args.node_target_type) + valset = HomogeneousDatasetAdapter(val_base, args.node_target_type) + testset = HomogeneousDatasetAdapter(test_base, args.node_target_type) + else: + basedir = os.path.join(datadir, f"{args.modelname}.pickle") + trainset = HomogeneousDatasetAdapter( + SimplePickleDataset(basedir=basedir, label="trainset", var_config=None), + args.node_target_type, + ) + valset = HomogeneousDatasetAdapter( + SimplePickleDataset(basedir=basedir, label="valset", var_config=None), + args.node_target_type, + ) + testset = HomogeneousDatasetAdapter( + SimplePickleDataset(basedir=basedir, label="testset", var_config=None), + args.node_target_type, + ) + + resolved_node_target_type = _resolve_node_target_type( + trainset[0], args.node_target_type + ) + if resolved_node_target_type != args.node_target_type: + info( + f"Resolved node_target_type '{args.node_target_type}' -> '{resolved_node_target_type}'" + ) + args.node_target_type = resolved_node_target_type + + # Sync input_node_features with the actual homogeneous data feature dim. + # After to_homogeneous(), x is zero-padded to the max feature dim across + # all node types, which is typically larger than the config's original + # input_node_features list. + actual_x_dim = trainset[0].x.shape[1] + voi = config["NeuralNetwork"]["Variables_of_interest"] + if len(voi["input_node_features"]) != actual_x_dim: + info( + f"Updating input_node_features: config has {len(voi['input_node_features'])} " + f"features but homogeneous data has {actual_x_dim} (zero-padded)." + ) + voi["input_node_features"] = list(range(actual_x_dim)) + + info( + "trainset,valset,testset size: %d %d %d" + % (len(trainset), len(valset), len(testset)) + ) + + (train_loader, val_loader, test_loader,) = hydragnn.preprocess.create_dataloaders( + trainset, valset, testset, config["NeuralNetwork"]["Training"]["batch_size"] + ) + + train_loader = HomogeneousBatchAdapter(train_loader) + val_loader = HomogeneousBatchAdapter(val_loader) + test_loader = HomogeneousBatchAdapter(test_loader) + + config = update_config(config, train_loader, val_loader, test_loader) + config = _to_jsonable(config) + hydragnn.utils.input_config_parsing.save_config(config, log_name) + + precision = config["NeuralNetwork"]["Training"].get("precision", "fp32") + model = hydragnn.models.create_model_config( + config=config["NeuralNetwork"], + verbosity=config["Verbosity"]["level"], + ) + + learning_rate = config["NeuralNetwork"]["Training"]["Optimizer"]["learning_rate"] + optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=5, min_lr=1e-5 + ) + + model, optimizer = hydragnn.utils.distributed.distributed_model_wrapper( + model, optimizer, config["Verbosity"]["level"] + ) + + _diag("Entering print_model") + print_model(model) + _diag("Exited print_model") + + _diag("Entering load_existing_model_config") + hydragnn.utils.model.load_existing_model_config( + model, config["NeuralNetwork"]["Training"], optimizer=optimizer + ) + _diag("Exited load_existing_model_config") + + _diag("Entering train_validate_test") + hydragnn.train.train_validate_test( + model, + optimizer, + train_loader, + val_loader, + test_loader, + writer, + scheduler, + config["NeuralNetwork"], + log_name, + config["Verbosity"]["level"], + create_plots=False, + precision=precision, + ) + _diag("Exited train_validate_test") + + hydragnn.utils.model.save_model(model, optimizer, log_name) + hydragnn.utils.profiling_and_tracing.print_timers(config["Verbosity"]["level"]) + if writer is not None: + writer.close() + + comm.Barrier() + if dist.is_initialized(): + dist.destroy_process_group() diff --git a/export_variables_and_run_interactive_job.sh b/export_variables_and_run_interactive_job.sh old mode 100755 new mode 100644 index a23680c9f..4af83f985 --- a/export_variables_and_run_interactive_job.sh +++ b/export_variables_and_run_interactive_job.sh @@ -9,17 +9,17 @@ function cmd() { time $@ } -HYDRAGNN_ROOT=/lustre/orion/lrn070/world-shared/mlupopa/Supercomputing2026/HydraGNN +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN # Load conda environemnt -source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm640.sh -source activate /lustre/orion/lrn070/world-shared/mlupopa/HydraGNN-Installation-Frontier/hydragnn_venv +source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm711.sh +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to HydragNN export PYTHONPATH=$PWD:$PYTHONPATH #export python path to use ADIOS2 v.2.10.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import adios2; print(adios2.__version__, adios2.__file__)" @@ -45,23 +45,38 @@ export HYDRAGNN_VALTEST=1 ## Getting error without these after 20 nodes export NCCL_P2P_LEVEL=NVL export NCCL_P2P_DISABLE=1 -export FI_MR_CACHE_MONITOR=disabled ## aws-ofi-rccl plugin settings -export TORCH_NCCL_HIGH_PRIORITY=1 -export FI_CXI_RDV_PROTO=alt_read +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib -export PATH_TO_THE_PLUGIN_DIRECTORY=/lustre/orion/lrn070/world-shared/mlupopa/AWI_OFI_RCCL_ROCm631/aws-ofi-rccl/lib -export LD_LIBRARY_PATH=${PATH_TO_THE_PLUGIN_DIRECTORY}:$LD_LIBRARY_PATH - export FI_MR_CACHE_MONITOR=kdreg2 # Required to avoid a deadlock. export FI_CXI_DEFAULT_CQ_SIZE=131072 # Ask the network stack to allocate additional space to process message completions. export FI_CXI_DEFAULT_TX_SIZE=2048 # Ask the network stack to allocate additional space to hold pending outgoing messages. export FI_CXI_RX_MATCH_MODE=hybrid # Allow the network stack to transition to software mode if necessary. - -export NCCL_NET_GDR_LEVEL=3 # Typically improves performance, but remove this setting if you encounter a hang/crash. +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" # Typically improves performance, but remove this setting if you encounter a hang/crash. export NCCL_CROSS_NIC=1 # On large systems, this NCCL setting has been found to improve performance export NCCL_SOCKET_IFNAME=hsn0 # NCCL/RCCL will use the high speed network to coordinate startup. +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 # Use high priority stream for the NCCL/RCCL Communicator. +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +# below are optional to debug RCCL stuff +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=INIT + +# The following have been found to help avoid hangs, but are not yet +# documented elsewhere +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 ## Checking diff --git a/export_variables_and_run_interactive_job_mptrj_example.sh b/export_variables_and_run_interactive_job_mptrj_example.sh old mode 100755 new mode 100644 index 49611daf1..b0242d9f3 --- a/export_variables_and_run_interactive_job_mptrj_example.sh +++ b/export_variables_and_run_interactive_job_mptrj_example.sh @@ -9,17 +9,17 @@ function cmd() { time $@ } -HYDRAGNN_ROOT=/lustre/orion/lrn070/world-shared/mlupopa/Supercomputing2026/HydraGNN +HYDRAGNN_ROOT=/lustre/orion/lrn078/proj-shared/HydraGNN # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm640.sh -source activate /lustre/orion/lrn070/world-shared/mlupopa/HydraGNN-Installation-Frontier/hydragnn_venv +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to HydragNN export PYTHONPATH=$PWD:$PYTHONPATH #export python path to use ADIOS2 v.2.10.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import adios2; print(adios2.__version__, adios2.__file__)" @@ -45,23 +45,38 @@ export HYDRAGNN_VALTEST=1 ## Getting error without these after 20 nodes export NCCL_P2P_LEVEL=NVL export NCCL_P2P_DISABLE=1 -export FI_MR_CACHE_MONITOR=disabled ## aws-ofi-rccl plugin settings -export TORCH_NCCL_HIGH_PRIORITY=1 -export FI_CXI_RDV_PROTO=alt_read +export PLUGIN_PATH=/ccs/sw/crusher/amdsw/aws-ofi-nccl/aws-ofi-nccl +export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:${PLUGIN_PATH}/lib -export PATH_TO_THE_PLUGIN_DIRECTORY=/lustre/orion/lrn070/world-shared/mlupopa/AWI_OFI_RCCL_ROCm631/aws-ofi-rccl/lib -export LD_LIBRARY_PATH=${PATH_TO_THE_PLUGIN_DIRECTORY}:$LD_LIBRARY_PATH - export FI_MR_CACHE_MONITOR=kdreg2 # Required to avoid a deadlock. export FI_CXI_DEFAULT_CQ_SIZE=131072 # Ask the network stack to allocate additional space to process message completions. export FI_CXI_DEFAULT_TX_SIZE=2048 # Ask the network stack to allocate additional space to hold pending outgoing messages. export FI_CXI_RX_MATCH_MODE=hybrid # Allow the network stack to transition to software mode if necessary. - -export NCCL_NET_GDR_LEVEL=3 # Typically improves performance, but remove this setting if you encounter a hang/crash. +export FI_CXI_RDV_PROTO=alt_read +export FI_CXI_DISABLE_HOST_REGISTER=1 + +export NCCL_NET_PLUGIN=${PLUGIN_PATH}/lib/librccl-net.so +export NCCL_NET_GDR_LEVEL="PHB" # Typically improves performance, but remove this setting if you encounter a hang/crash. export NCCL_CROSS_NIC=1 # On large systems, this NCCL setting has been found to improve performance export NCCL_SOCKET_IFNAME=hsn0 # NCCL/RCCL will use the high speed network to coordinate startup. +export NCCL_NET="AWS Libfabric" + +export TORCH_NCCL_HIGH_PRIORITY=1 # Use high priority stream for the NCCL/RCCL Communicator. +export GPU_MAX_HW_QUEUES=2 + +export HSA_FORCE_FINE_GRAIN_PCIE=1 + +# below are optional to debug RCCL stuff +# export NCCL_DEBUG=INFO +# export NCCL_DEBUG_SUBSYS=INIT + +# The following have been found to help avoid hangs, but are not yet +# documented elsewhere +export FI_CXI_RDZV_EAGER_SIZE=0 +export FI_CXI_RDZV_GET_MIN=0 +export FI_CXI_RDZV_THRESHOLD=0 ## Checking diff --git a/hydragnn/models/__init__.py b/hydragnn/models/__init__.py index 6997e4461..404678176 100644 --- a/hydragnn/models/__init__.py +++ b/hydragnn/models/__init__.py @@ -3,5 +3,15 @@ from .GINStack import GINStack from .PNAStack import PNAStack from .GINStack import GINStack +from .heterogeneous import ( + HeteroBase, + HeteroGINStack, + HeteroSAGEStack, + HeteroGATStack, + HeteroPNAStack, + HeteroHGTStack, + HeteroHEATStack, + HeteroRGATStack, +) from .create import create_model, create_model_config from .MultiTaskModelMP import MultiTaskModelMP, DualOptimizer diff --git a/hydragnn/models/create.py b/hydragnn/models/create.py index 72e24e3e5..00d84a582 100644 --- a/hydragnn/models/create.py +++ b/hydragnn/models/create.py @@ -12,7 +12,7 @@ import os import torch from torch_geometric.data import Data -from typing import List, Union +from typing import Dict, List, Union import torch_scatter @@ -30,6 +30,15 @@ from hydragnn.models.PNAEqStack import PNAEqStack from hydragnn.models.PAINNStack import PAINNStack from hydragnn.models.MACEStack import MACEStack +from hydragnn.models.heterogeneous import ( + HeteroGINStack, + HeteroSAGEStack, + HeteroGATStack, + HeteroPNAStack, + HeteroHGTStack, + HeteroHEATStack, + HeteroRGATStack, +) # InteratomicPotential functionality is now implemented via wrapper composition @@ -42,6 +51,8 @@ def create_model_config( config: dict, verbosity: int = 0, use_gpu: bool = True, + metadata=None, + node_input_dims=None, ): model = create_model( mpnn_type=config["Architecture"]["mpnn_type"], @@ -96,6 +107,23 @@ def create_model_config( "graph_attr_conditioning_mode", "concat_node" ), graph_pooling=config["Architecture"].get("graph_pooling", "mean"), + hetero_pooling_mode=config["Architecture"].get("hetero_pooling_mode", "sum"), + node_target_type=config["Architecture"].get("node_target_type", None), + share_relation_weights=config["Architecture"].get( + "share_relation_weights", False + ), + metadata=metadata, + node_input_dims=node_input_dims, + hetero_attention_heads=config["Architecture"].get("hetero_attention_heads", 4), + hetero_attention_negative_slope=config["Architecture"].get( + "hetero_attention_negative_slope", 0.2 + ), + hetero_edge_type_emb_dim=config["Architecture"].get( + "hetero_edge_type_emb_dim", 16 + ), + hetero_edge_attr_emb_dim=config["Architecture"].get( + "hetero_edge_attr_emb_dim", 16 + ), verbosity=verbosity, use_gpu=use_gpu, ) @@ -128,7 +156,7 @@ def create_model( initial_bias: float = None, num_nodes: int = None, max_neighbours: int = None, - edge_dim: int = None, + edge_dim: Union[int, Dict[str, int], None] = None, pna_deg: torch.tensor = None, num_before_skip: int = None, num_after_skip: int = None, @@ -156,6 +184,15 @@ def create_model( use_graph_attr_conditioning: bool = False, graph_attr_conditioning_mode: str = "fuse_pool", graph_pooling: str = "mean", + hetero_pooling_mode: str = "sum", + node_target_type: str = None, + share_relation_weights: bool = False, + metadata=None, + node_input_dims=None, + hetero_attention_heads: int = 4, + hetero_attention_negative_slope: float = 0.2, + hetero_edge_type_emb_dim: int = 16, + hetero_edge_attr_emb_dim: int = 16, verbosity: int = 0, use_gpu: bool = True, ): @@ -580,6 +617,221 @@ def create_model( use_graph_attr_conditioning=use_graph_attr_conditioning, graph_attr_conditioning_mode=graph_attr_conditioning_mode, ) + + elif mpnn_type == "HeteroGIN": + model = HeteroGINStack( + input_dim, + hidden_dim, + output_dim, + pe_dim, + global_attn_engine, + global_attn_type, + global_attn_heads, + output_type, + output_heads, + activation_function, + loss_function_type, + equivariance, + loss_weights=task_weights, + freeze_conv=freeze_conv, + initial_bias=initial_bias, + num_conv_layers=num_conv_layers, + num_nodes=num_nodes, + graph_pooling=graph_pooling, + use_graph_attr_conditioning=use_graph_attr_conditioning, + graph_attr_conditioning_mode=graph_attr_conditioning_mode, + hetero_pooling_mode=hetero_pooling_mode, + node_target_type=node_target_type, + share_relation_weights=share_relation_weights, + metadata=metadata, + node_input_dims=node_input_dims, + ) + + elif mpnn_type == "HeteroSAGE": + model = HeteroSAGEStack( + input_dim, + hidden_dim, + output_dim, + pe_dim, + global_attn_engine, + global_attn_type, + global_attn_heads, + output_type, + output_heads, + activation_function, + loss_function_type, + equivariance, + loss_weights=task_weights, + freeze_conv=freeze_conv, + initial_bias=initial_bias, + num_conv_layers=num_conv_layers, + num_nodes=num_nodes, + graph_pooling=graph_pooling, + use_graph_attr_conditioning=use_graph_attr_conditioning, + graph_attr_conditioning_mode=graph_attr_conditioning_mode, + hetero_pooling_mode=hetero_pooling_mode, + node_target_type=node_target_type, + share_relation_weights=share_relation_weights, + metadata=metadata, + node_input_dims=node_input_dims, + ) + + elif mpnn_type == "HeteroGAT": + heads = 6 + negative_slope = 0.05 + model = HeteroGATStack( + heads, + negative_slope, + edge_dim, + input_dim, + hidden_dim, + output_dim, + pe_dim, + global_attn_engine, + global_attn_type, + global_attn_heads, + output_type, + output_heads, + activation_function, + loss_function_type, + equivariance, + loss_weights=task_weights, + freeze_conv=freeze_conv, + initial_bias=initial_bias, + num_conv_layers=num_conv_layers, + num_nodes=num_nodes, + graph_pooling=graph_pooling, + use_graph_attr_conditioning=use_graph_attr_conditioning, + graph_attr_conditioning_mode=graph_attr_conditioning_mode, + hetero_pooling_mode=hetero_pooling_mode, + node_target_type=node_target_type, + share_relation_weights=share_relation_weights, + metadata=metadata, + node_input_dims=node_input_dims, + ) + + elif mpnn_type == "HeteroPNA": + assert pna_deg is not None, "HeteroPNA requires degree input." + model = HeteroPNAStack( + pna_deg, + edge_dim, + input_dim, + hidden_dim, + output_dim, + pe_dim, + global_attn_engine, + global_attn_type, + global_attn_heads, + output_type, + output_heads, + activation_function, + loss_function_type, + equivariance, + loss_weights=task_weights, + freeze_conv=freeze_conv, + initial_bias=initial_bias, + num_conv_layers=num_conv_layers, + num_nodes=num_nodes, + graph_pooling=graph_pooling, + use_graph_attr_conditioning=use_graph_attr_conditioning, + graph_attr_conditioning_mode=graph_attr_conditioning_mode, + hetero_pooling_mode=hetero_pooling_mode, + node_target_type=node_target_type, + share_relation_weights=share_relation_weights, + metadata=metadata, + node_input_dims=node_input_dims, + ) + elif mpnn_type == "HeteroRGAT": + model = HeteroRGATStack( + hetero_attention_heads, + hetero_attention_negative_slope, + edge_dim, + input_dim, + hidden_dim, + output_dim, + pe_dim, + global_attn_engine, + global_attn_type, + global_attn_heads, + output_type, + output_heads, + activation_function, + loss_function_type, + equivariance, + loss_weights=task_weights, + freeze_conv=freeze_conv, + initial_bias=initial_bias, + num_conv_layers=num_conv_layers, + num_nodes=num_nodes, + graph_pooling=graph_pooling, + use_graph_attr_conditioning=use_graph_attr_conditioning, + graph_attr_conditioning_mode=graph_attr_conditioning_mode, + hetero_pooling_mode=hetero_pooling_mode, + node_target_type=node_target_type, + share_relation_weights=share_relation_weights, + metadata=metadata, + node_input_dims=node_input_dims, + ) + elif mpnn_type == "HeteroHGT": + model = HeteroHGTStack( + hetero_attention_heads, + input_dim, + hidden_dim, + output_dim, + pe_dim, + global_attn_engine, + global_attn_type, + global_attn_heads, + output_type, + output_heads, + activation_function, + loss_function_type, + equivariance, + loss_weights=task_weights, + freeze_conv=freeze_conv, + initial_bias=initial_bias, + num_conv_layers=num_conv_layers, + num_nodes=num_nodes, + graph_pooling=graph_pooling, + use_graph_attr_conditioning=use_graph_attr_conditioning, + graph_attr_conditioning_mode=graph_attr_conditioning_mode, + hetero_pooling_mode=hetero_pooling_mode, + node_target_type=node_target_type, + share_relation_weights=share_relation_weights, + metadata=metadata, + node_input_dims=node_input_dims, + ) + elif mpnn_type == "HeteroHEAT": + model = HeteroHEATStack( + hetero_attention_heads, + hetero_edge_type_emb_dim, + hetero_edge_attr_emb_dim, + input_dim, + hidden_dim, + output_dim, + pe_dim, + global_attn_engine, + global_attn_type, + global_attn_heads, + output_type, + output_heads, + activation_function, + loss_function_type, + equivariance, + loss_weights=task_weights, + freeze_conv=freeze_conv, + initial_bias=initial_bias, + num_conv_layers=num_conv_layers, + num_nodes=num_nodes, + graph_pooling=graph_pooling, + use_graph_attr_conditioning=use_graph_attr_conditioning, + graph_attr_conditioning_mode=graph_attr_conditioning_mode, + hetero_pooling_mode=hetero_pooling_mode, + node_target_type=node_target_type, + share_relation_weights=share_relation_weights, + metadata=metadata, + node_input_dims=node_input_dims, + ) else: raise ValueError("Unknown mpnn_type: {0}".format(mpnn_type)) diff --git a/hydragnn/models/heterogeneous/HeteroBase.py b/hydragnn/models/heterogeneous/HeteroBase.py new file mode 100644 index 000000000..bbe3d6094 --- /dev/null +++ b/hydragnn/models/heterogeneous/HeteroBase.py @@ -0,0 +1,897 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +import torch +from torch.nn import Module, ModuleList, ModuleDict, Linear, Sequential +from torch_geometric.nn import ( + BatchNorm, + HeteroConv, + global_add_pool, + global_max_pool, + global_mean_pool, +) + +from hydragnn.utils.model import activation_function_selection, loss_function_selection +from hydragnn.utils.distributed import get_device +from hydragnn.models.Base import MLPNode + + +class HeteroBase(Module): + """Base class for heterogeneous message passing models. + + This reuses HydraGNN's multi-head decoding logic while allowing hetero + message passing via PyG HeteroConv. + """ + + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: list, + pe_dim: int, + global_attn_engine: str, + global_attn_type: str, + global_attn_heads: int, + output_type: list, + config_heads: dict, + activation_function_type: str, + loss_function_type: str, + equivariance: bool, + ilossweights_hyperp: int = 1, + loss_weights: list = None, + ilossweights_nll: int = 0, + freeze_conv: bool = False, + initial_bias=None, + dropout: float = 0.25, + num_conv_layers: int = 16, + num_nodes: int = None, + graph_pooling: str = "mean", + use_graph_attr_conditioning: bool = False, + graph_attr_conditioning_mode: str = "concat_node", + hetero_pooling_mode: str = "sum", + node_target_type: str = None, + share_relation_weights: bool = False, + node_input_dims: dict | None = None, + metadata=None, + ): + super().__init__() + + if global_attn_engine: + raise NotImplementedError( + "HeteroBase does not yet support global attention. Set global_attn_engine=None." + ) + + self.device = get_device() + self.input_dim = input_dim + self.pe_dim = pe_dim + self.hidden_dim = hidden_dim + self.dropout = dropout + self.num_conv_layers = num_conv_layers + self.num_nodes = num_nodes + self.graph_convs = ModuleList() + self.feature_layers = ModuleList() + self.node_embedders = ModuleDict() + self._node_input_dims = node_input_dims + self.node_target_type = node_target_type + self.share_relation_weights = share_relation_weights + self._metadata = metadata + self._initialized = False + self._pending_node_conv_init = False + self._node_conv_head_specs = [] + + self.global_attn_engine = global_attn_engine + self.global_attn_type = global_attn_type + self.global_attn_heads = global_attn_heads + + self.heads_NN = ModuleList() + self.config_heads = config_heads + self.head_type = output_type + self.head_dims = output_dim + self.num_heads = len(self.head_dims) + self.convs_node_hidden = ModuleDict({}) + self.batch_norms_node_hidden = ModuleDict({}) + self.convs_node_output = ModuleDict({}) + self.batch_norms_node_output = ModuleDict({}) + + self.equivariance = equivariance + self.activation_function = activation_function_selection( + activation_function_type + ) + + self.use_graph_attr_conditioning = use_graph_attr_conditioning + self.graph_attr_conditioning_mode = graph_attr_conditioning_mode.lower() + if self.graph_attr_conditioning_mode not in ( + "film", + "concat_node", + "fuse_pool", + ): + raise ValueError( + "graph_attr_conditioning_mode must be one of: 'film', 'concat_node', 'fuse_pool'." + ) + + # output variance for Gaussian negative log likelihood loss + self.var_output = 0 + if loss_function_type == "GaussianNLLLoss": + self.var_output = 1 + self.loss_function_type = loss_function_type + self.loss_function = loss_function_selection(loss_function_type) + self.ilossweights_nll = ilossweights_nll + self.ilossweights_hyperp = ilossweights_hyperp + + if loss_weights is None: + loss_weights = [1.0] * self.num_heads + + if self.ilossweights_hyperp * self.ilossweights_nll == 1: + raise ValueError( + "ilossweights_hyperp and ilossweights_nll cannot be both set to 1." + ) + if self.ilossweights_hyperp == 1: + if len(loss_weights) != self.num_heads: + raise ValueError( + "Inconsistent number of loss weights and tasks: " + + str(len(loss_weights)) + + " VS " + + str(self.num_heads) + ) + else: + self.loss_weights = loss_weights + weightabssum = sum(abs(number) for number in self.loss_weights) + self.loss_weights = [iw / weightabssum for iw in self.loss_weights] + + # Graph pooling policy + pool_mode = graph_pooling.lower() + if pool_mode == "sum": + pool_mode = "add" + pool_map = { + "mean": (global_mean_pool, "mean"), + "add": (global_add_pool, "sum"), + "max": (global_max_pool, "max"), + } + if pool_mode not in pool_map: + raise ValueError("Unsupported graph_pooling: " + graph_pooling) + self.graph_pooling = pool_mode + self.graph_pool_fn, self.graph_pool_reduction = pool_map[pool_mode] + + if hetero_pooling_mode not in ("sum", "mean"): + raise ValueError("hetero_pooling_mode must be 'sum' or 'mean'.") + self.hetero_pooling_mode = hetero_pooling_mode + + def _pool_graph_features(x_tensor, batch_tensor): + if batch_tensor is None: + if self.graph_pool_reduction == "mean": + return x_tensor.mean(dim=0, keepdim=True) + if self.graph_pool_reduction == "max": + return x_tensor.max(dim=0, keepdim=True).values + return x_tensor.sum(dim=0, keepdim=True) + return self.graph_pool_fn(x_tensor, batch_tensor.to(x_tensor.device)) + + self._pool_graph_features = _pool_graph_features + + self.freeze_conv = freeze_conv + self.initial_bias = initial_bias + + # Graph conditioning modules (lazy) + self.graph_conditioner = None + self.graph_concat_projector = None + self.graph_concat_projector_in_dim = None + self.graph_pool_projector = None + self.graph_pool_projector_in_dim = None + + self._multihead() + if self.initial_bias is not None: + self._set_bias() + + self.conv_checkpointing = False + + if self._metadata is not None: + self._init_conv() + + if self._node_input_dims: + self._init_node_embedders_from_dims(self._node_input_dims) + + def _init_node_embedders_from_dims(self, node_input_dims): + for node_type, in_dim in node_input_dims.items(): + if node_type not in self.node_embedders: + self.node_embedders[node_type] = Linear(int(in_dim), self.hidden_dim) + if self.node_embedders[node_type].weight.device != self.device: + self.node_embedders[node_type] = self.node_embedders[node_type].to( + self.device + ) + + def _ensure_node_embedders(self, x_dict): + for node_type, x in x_dict.items(): + if node_type not in self.node_embedders: + self.node_embedders[node_type] = Linear(x.size(-1), self.hidden_dim) + if self.node_embedders[node_type].weight.device != x.device: + self.node_embedders[node_type] = self.node_embedders[node_type].to( + x.device + ) + + def _maybe_init_metadata(self, data): + if self._metadata is None: + self._metadata = data.metadata() + if not self._initialized: + self._init_conv() + # Move lazily-created conv modules to the same device as the model. + device = next(self.parameters()).device + self.graph_convs.to(device) + self.feature_layers.to(device) + if self._pending_node_conv_init: + self._init_node_conv() + self._finalize_node_conv_heads() + self._pending_node_conv_init = False + + def _build_hetero_conv(self, input_dim: int, output_dim: int): + conv_dict = {} + shared_conv = None + for edge_type in self._metadata[1]: + edge_dim = self._resolve_edge_dim_for_type(edge_type) + if self.share_relation_weights: + if shared_conv is None: + shared_conv = self.get_conv( + input_dim, output_dim, edge_dim=edge_dim + ) + conv_dict[edge_type] = shared_conv + else: + conv_dict[edge_type] = self.get_conv( + input_dim, output_dim, edge_dim=edge_dim + ) + return HeteroConv(conv_dict, aggr="sum") + + def _build_hetero_conv_node_head(self, input_dim: int, output_dim: int): + conv_dict = {} + shared_conv = None + for edge_type in self._metadata[1]: + edge_dim = self._resolve_edge_dim_for_type(edge_type) + if self.share_relation_weights: + if shared_conv is None: + shared_conv = self.get_conv( + input_dim, output_dim, edge_dim=edge_dim + ) + conv_dict[edge_type] = shared_conv + else: + conv_dict[edge_type] = self.get_conv( + input_dim, output_dim, edge_dim=edge_dim + ) + return HeteroConv(conv_dict, aggr="sum") + + def _resolve_edge_dim_for_type(self, edge_type): + """Return the edge_dim for a specific edge type. + + When ``self.edge_dim`` is a dict mapping relation names to widths, + look up the relation (middle element of the triple). Returns ``None`` + for featureless edge types. When ``self.edge_dim`` is an int (or + absent), return it unchanged for all edge types. + """ + edge_dim = getattr(self, "edge_dim", None) + if isinstance(edge_dim, dict): + _, rel, _ = edge_type + return edge_dim.get(rel) + return edge_dim + + def _init_conv(self): + self.graph_convs = ModuleList() + self.feature_layers = ModuleList() + for layer_idx in range(self.num_conv_layers): + in_dim = self.hidden_dim if layer_idx > 0 else self.hidden_dim + out_dim = self.hidden_dim + self.graph_convs.append(self._build_hetero_conv(in_dim, out_dim)) + node_norms = ModuleDict({}) + for node_type in self._metadata[0]: + node_norms[node_type] = BatchNorm(out_dim) + self.feature_layers.append(node_norms) + self._initialized = True + + def _init_node_conv(self): + nodeconfiglist = self.config_heads["node"] + assert ( + self.num_branches == len(nodeconfiglist) or self.num_branches == 1 + ), "assuming node head has the same branches as graph head, if any" + for branchdict in nodeconfiglist: + if branchdict["architecture"]["type"] != "conv": + return + + node_feature_ind = [ + i for i, head_type in enumerate(self.head_type) if head_type == "node" + ] + if len(node_feature_ind) == 0: + return + + for branchdict in nodeconfiglist: + branchtype = branchdict["type"] + brancharct = branchdict["architecture"] + num_conv_layers_node = brancharct["num_headlayers"] + hidden_dim_node = brancharct["dim_headlayers"] + + convs_node_hidden = ModuleList() + batch_norms_node_hidden = ModuleList() + convs_node_output = ModuleList() + batch_norms_node_output = ModuleList() + + convs_node_hidden.append( + self._build_hetero_conv_node_head(self.hidden_dim, hidden_dim_node[0]) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm(hidden_dim_node[0]) + batch_norms_node_hidden.append(bn_dict) + + for ilayer in range(num_conv_layers_node - 1): + convs_node_hidden.append( + self._build_hetero_conv_node_head( + hidden_dim_node[ilayer], hidden_dim_node[ilayer + 1] + ) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm(hidden_dim_node[ilayer + 1]) + batch_norms_node_hidden.append(bn_dict) + + for ihead in node_feature_ind: + convs_node_output.append( + self._build_hetero_conv_node_head( + hidden_dim_node[-1], + self.head_dims[ihead] * (1 + self.var_output), + ) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm( + self.head_dims[ihead] * (1 + self.var_output) + ) + batch_norms_node_output.append(bn_dict) + + self.convs_node_hidden[branchtype] = convs_node_hidden + self.batch_norms_node_hidden[branchtype] = batch_norms_node_hidden + self.convs_node_output[branchtype] = convs_node_output + self.batch_norms_node_output[branchtype] = batch_norms_node_output + + def _freeze_conv(self): + for module in [self.graph_convs, self.feature_layers]: + for layer in module: + for param in layer.parameters(): + param.requires_grad = False + + def _set_bias(self): + for head, type in zip(self.heads_NN, self.head_type): + if type == "graph": + head[-1].bias.data.fill_(self.initial_bias) + + def _multihead(self): + self.graph_shared = ModuleDict({}) + dim_sharedlayers = 0 + self.num_branches = 1 + if "graph" in self.config_heads: + self.num_branches = len(self.config_heads["graph"]) + for branchdict in self.config_heads["graph"]: + denselayers = [] + dim_sharedlayers = branchdict["architecture"]["dim_sharedlayers"] + denselayers.append(Linear(self.hidden_dim, dim_sharedlayers)) + denselayers.append(self.activation_function) + for _ in range(branchdict["architecture"]["num_sharedlayers"] - 1): + denselayers.append(Linear(dim_sharedlayers, dim_sharedlayers)) + denselayers.append(self.activation_function) + self.graph_shared[branchdict["type"]] = Sequential(*denselayers) + + if "node" in self.config_heads: + if self._metadata is None: + self._pending_node_conv_init = True + else: + self._init_node_conv() + + inode_feature = 0 + for ihead in range(self.num_heads): + head_NN = ModuleDict({}) + if self.head_type[ihead] == "graph": + for branchdict in self.config_heads["graph"]: + branchtype = branchdict["type"] + brancharct = branchdict["architecture"] + dim_sharedlayers = brancharct["dim_sharedlayers"] + num_head_hidden = brancharct["num_headlayers"] + dim_head_hidden = brancharct["dim_headlayers"] + denselayers = [] + denselayers.append(Linear(dim_sharedlayers, dim_head_hidden[0])) + denselayers.append(self.activation_function) + for ilayer in range(num_head_hidden - 1): + denselayers.append( + Linear(dim_head_hidden[ilayer], dim_head_hidden[ilayer + 1]) + ) + denselayers.append(self.activation_function) + denselayers.append( + Linear( + dim_head_hidden[-1], + self.head_dims[ihead] * (1 + self.var_output), + ) + ) + head_NN[branchtype] = Sequential(*denselayers) + elif self.head_type[ihead] == "node": + for branchdict in self.config_heads["node"]: + branchtype = branchdict["type"] + brancharct = branchdict["architecture"] + hidden_dim_node = brancharct["dim_headlayers"] + node_NN_type = brancharct["type"] + if node_NN_type == "mlp" or node_NN_type == "mlp_per_node": + self.num_mlp = 1 if node_NN_type == "mlp" else self.num_nodes + if node_NN_type == "mlp_per_node": + assert ( + self.num_nodes is not None + ), "num_nodes must be provided for mlp_per_node; use 'mlp' for variable-size graphs" + head_NN[branchtype] = MLPNode( + self.hidden_dim, + self.head_dims[ihead] * (1 + self.var_output), + self.num_mlp, + hidden_dim_node, + node_NN_type, + self.activation_function, + num_nodes=self.num_nodes + if node_NN_type == "mlp_per_node" + else None, + ) + elif node_NN_type == "conv": + head_NN[branchtype] = ModuleList() + if self._metadata is None: + self._node_conv_head_specs.append( + (ihead, branchtype, inode_feature) + ) + inode_feature += 1 + else: + for conv, batch_norm in zip( + self.convs_node_hidden[branchtype], + self.batch_norms_node_hidden[branchtype], + ): + head_NN[branchtype].append(conv) + head_NN[branchtype].append(batch_norm) + head_NN[branchtype].append( + self.convs_node_output[branchtype][inode_feature] + ) + head_NN[branchtype].append( + self.batch_norms_node_output[branchtype][inode_feature] + ) + inode_feature += 1 + else: + raise ValueError( + "HeteroBase only supports node heads with 'mlp', 'mlp_per_node', or 'conv'." + ) + else: + raise ValueError( + "Unknown head type" + + self.head_type[ihead] + + "; currently only support 'graph' or 'node'" + ) + self.heads_NN.append(head_NN) + + def _get_batch_dict(self, data, x_dict): + batch_dict = None + try: + batch_dict = data.batch_dict + except (AttributeError, KeyError): + batch_dict = None + if batch_dict is not None: + return batch_dict + batch_dict = {} + for node_type, x in x_dict.items(): + batch_dict[node_type] = torch.zeros( + x.size(0), device=x.device, dtype=torch.long + ) + return batch_dict + + def _get_edge_attr_dict(self, data): + if not getattr(self, "is_edge_model", False): + return None + edge_attr_dict = None + try: + edge_attr_dict = data.edge_attr_dict + except (AttributeError, KeyError): + edge_attr_dict = None + return edge_attr_dict + + def _pool_hetero_graph_features(self, x_dict, batch_dict): + pooled = [] + for node_type, x in x_dict.items(): + pooled.append(self._pool_graph_features(x, batch_dict[node_type])) + if len(pooled) == 1: + return pooled[0] + if self.hetero_pooling_mode == "sum": + return torch.stack(pooled, dim=0).sum(dim=0) + return torch.stack(pooled, dim=0).mean(dim=0) + + def _finalize_node_conv_heads(self): + if not self._node_conv_head_specs: + return + for head_index, branchtype, output_index in self._node_conv_head_specs: + headloc = self.heads_NN[head_index] + if branchtype not in headloc: + headloc[branchtype] = ModuleList() + for conv, batch_norm in zip( + self.convs_node_hidden[branchtype], + self.batch_norms_node_hidden[branchtype], + ): + headloc[branchtype].append(conv) + headloc[branchtype].append(batch_norm) + headloc[branchtype].append(self.convs_node_output[branchtype][output_index]) + headloc[branchtype].append( + self.batch_norms_node_output[branchtype][output_index] + ) + + def _ensure_graph_conditioner(self, graph_attr_dim: int, device): + if self.graph_conditioner is None: + hidden = max(self.hidden_dim, graph_attr_dim) + self.graph_conditioner = Sequential( + Linear(graph_attr_dim, hidden), + self.activation_function, + Linear(hidden, 2 * self.hidden_dim), + ) + if self.graph_conditioner[0].weight.device != device: + self.graph_conditioner = self.graph_conditioner.to(device) + + def _ensure_graph_concat_projector( + self, graph_attr_dim: int, channel_dim: int, device + ): + in_dim = channel_dim + graph_attr_dim + if (self.graph_concat_projector is None) or ( + self.graph_concat_projector_in_dim != in_dim + ): + self.graph_concat_projector = Linear(in_dim, channel_dim) + self.graph_concat_projector_in_dim = in_dim + if self.graph_concat_projector.weight.device != device: + self.graph_concat_projector = self.graph_concat_projector.to(device) + + def _ensure_graph_pool_projector( + self, graph_attr_dim: int, channel_dim: int, device + ): + in_dim = channel_dim + graph_attr_dim + if (self.graph_pool_projector is None) or ( + self.graph_pool_projector_in_dim != in_dim + ): + self.graph_pool_projector = Sequential( + Linear(in_dim, channel_dim), + self.activation_function, + Linear(channel_dim, channel_dim), + ) + self.graph_pool_projector_in_dim = in_dim + if self.graph_pool_projector[0].weight.device != device: + self.graph_pool_projector = self.graph_pool_projector.to(device) + + def _apply_graph_conditioning(self, inv_node_feat, batch, data): + if not self.use_graph_attr_conditioning: + return inv_node_feat + + if not hasattr(data, "graph_attr") or data.graph_attr is None: + raise ValueError( + "use_graph_attr_conditioning=True but data.graph_attr is missing." + ) + + graph_attr = data.graph_attr + graph_attr = graph_attr.to(inv_node_feat.device).float() + + if batch is None: + batch = torch.zeros( + inv_node_feat.size(0), device=inv_node_feat.device, dtype=torch.long + ) + + num_graphs = int(batch.max().item() + 1) + + if graph_attr.dim() == 1: + if graph_attr.numel() % num_graphs == 0: + feat_dim = graph_attr.numel() // num_graphs + graph_attr = graph_attr.view(num_graphs, feat_dim) + else: + raise ValueError( + f"One-dimensional graph_attr with numel={graph_attr.numel()} is not divisible by num_graphs={num_graphs}." + ) + elif graph_attr.dim() == 2: + if graph_attr.size(0) != num_graphs: + raise ValueError( + f"graph_attr first dim {graph_attr.size(0)} does not match num_graphs={num_graphs}." + ) + else: + raise ValueError( + f"Unsupported graph_attr ndim={graph_attr.dim()}; expected 1/2." + ) + + if self.graph_attr_conditioning_mode == "film": + self._ensure_graph_conditioner(graph_attr.size(-1), inv_node_feat.device) + + scale_shift = self.graph_conditioner(graph_attr) + scale, shift = scale_shift.split(self.hidden_dim, dim=-1) + scale = torch.tanh(scale) + + channel_dim = inv_node_feat.size(-1) + scale_b = scale[batch] + shift_b = shift[batch] + if channel_dim != self.hidden_dim: + if channel_dim % self.hidden_dim != 0: + raise ValueError( + f"Graph conditioning expects channels divisible by hidden_dim (got {channel_dim} vs {self.hidden_dim})." + ) + factor = channel_dim // self.hidden_dim + scale_b = scale_b.repeat_interleave(factor, dim=-1) + shift_b = shift_b.repeat_interleave(factor, dim=-1) + + return inv_node_feat * (1 + scale_b) + shift_b + + if self.graph_attr_conditioning_mode == "concat_node": + channel_dim = inv_node_feat.size(-1) + self._ensure_graph_concat_projector( + graph_attr_dim=graph_attr.size(-1), + channel_dim=channel_dim, + device=inv_node_feat.device, + ) + graph_attr_b = graph_attr[batch] + fused = torch.cat([inv_node_feat, graph_attr_b], dim=-1) + return self.graph_concat_projector(fused) + + if self.graph_attr_conditioning_mode == "fuse_pool": + return inv_node_feat + + raise ValueError( + f"Unsupported graph_attr_conditioning_mode: {self.graph_attr_conditioning_mode}" + ) + + def _apply_graph_pool_conditioning(self, x_graph, data): + if not self.use_graph_attr_conditioning: + return x_graph + if self.graph_attr_conditioning_mode != "fuse_pool": + return x_graph + if not hasattr(data, "graph_attr") or data.graph_attr is None: + raise ValueError( + "use_graph_attr_conditioning=True but data.graph_attr is missing." + ) + + graph_attr = data.graph_attr + num_graphs = x_graph.size(0) + + if graph_attr.dim() == 1: + if graph_attr.numel() % num_graphs == 0: + feat_dim = graph_attr.numel() // num_graphs + graph_attr = graph_attr.view(num_graphs, feat_dim) + else: + raise ValueError( + f"One-dimensional graph attribute with graph_attr.numel()={graph_attr.numel()} is not divisible by num_graphs={num_graphs}." + ) + elif graph_attr.dim() == 2: + if graph_attr.size(0) != num_graphs: + raise ValueError( + f"graph_attr batch size does not match pooled graph embeddings: graph_attr={tuple(graph_attr.size())}, num_graphs={num_graphs}" + ) + else: + raise ValueError( + f"Unsupported graph_attr ndim={graph_attr.dim()}; expected 1/2." + ) + + graph_attr = graph_attr.to(x_graph.device).float() + + self._ensure_graph_pool_projector( + graph_attr_dim=graph_attr.size(-1), + channel_dim=x_graph.size(-1), + device=x_graph.device, + ) + + if graph_attr.size(0) != num_graphs: + raise ValueError( + f"graph_attr batch size does not match pooled graph embeddings: " + f"graph_attr={tuple(graph_attr.size())}, x_graph={tuple(x_graph.size())}, num_graphs={num_graphs}" + ) + + fused = torch.cat([x_graph, graph_attr], dim=-1) + return self.graph_pool_projector(fused) + + def forward(self, data): + self._maybe_init_metadata(data) + + device = next(self.parameters()).device + if hasattr(data, "to"): + data = data.to(device) + + if hasattr(data, "node_types"): + for node_type in data.node_types: + store = data[node_type] + if hasattr(store, "x") and store.x is not None: + store.x = store.x.to(device) + if hasattr(data, "edge_types"): + for edge_type in data.edge_types: + store = data[edge_type] + if hasattr(store, "edge_index") and store.edge_index is not None: + store.edge_index = store.edge_index.to(device) + if hasattr(store, "edge_attr") and store.edge_attr is not None: + store.edge_attr = store.edge_attr.to(device) + + x_dict = {node_type: x.to(device) for node_type, x in data.x_dict.items()} + self._ensure_node_embedders(x_dict) + x_dict = { + node_type: self.node_embedders[node_type](x.float()) + for node_type, x in x_dict.items() + } + + batch_dict = self._get_batch_dict(data, x_dict) + edge_attr_dict = self._get_edge_attr_dict(data) + + for conv, node_norms in zip(self.graph_convs, self.feature_layers): + if edge_attr_dict is None: + x_dict = conv(x_dict, data.edge_index_dict) + else: + x_dict = conv(x_dict, data.edge_index_dict, edge_attr_dict) + for node_type, x in x_dict.items(): + x = self._apply_graph_conditioning(x, batch_dict[node_type], data) + x = node_norms[node_type](x) + x = self.activation_function(x) + x_dict[node_type] = x + return self._decode_from_x_dict(x_dict, batch_dict, data, edge_attr_dict) + + def _decode_from_x_dict(self, x_dict, batch_dict, data, edge_attr_dict): + x_graph = self._pool_hetero_graph_features(x_dict, batch_dict) + x_graph = self._apply_graph_pool_conditioning(x_graph, data) + + # Prepare dataset_name for multi-branch heads + if not hasattr(data, "dataset_name"): + num_graphs = x_graph.size(0) + data.dataset_name = torch.zeros( + (num_graphs, 1), device=x_graph.device, dtype=torch.long + ) + else: + data.dataset_name = data.dataset_name.to(x_graph.device) + + outputs = [] + outputs_var = [] + + datasetIDs = data.dataset_name.unique() + + for head_dim, headloc, type_head in zip( + self.head_dims, self.heads_NN, self.head_type + ): + if type_head == "graph": + head = torch.zeros( + (len(data.dataset_name), head_dim), device=x_graph.device + ) + headvar = torch.zeros( + (len(data.dataset_name), head_dim * self.var_output), + device=x_graph.device, + ) + if self.num_branches == 1: + head_device = next( + self.graph_shared["branch-0"].parameters() + ).device + x_graph = x_graph.to(head_device) + x_graph_head = self.graph_shared["branch-0"](x_graph) + output_head = headloc["branch-0"](x_graph_head) + head = output_head[:, :head_dim] + headvar = output_head[:, head_dim:] ** 2 + else: + for ID in datasetIDs: + mask = data.dataset_name == ID + mask = mask[:, 0] + branchtype = f"branch-{ID.item()}" + head_device = next( + self.graph_shared[branchtype].parameters() + ).device + x_graph = x_graph.to(head_device) + x_graph_head = self.graph_shared[branchtype](x_graph[mask, :]) + output_head = headloc[branchtype](x_graph_head) + head[mask] = output_head[:, :head_dim] + headvar[mask] = output_head[:, head_dim:] ** 2 + outputs.append(head) + outputs_var.append(headvar) + else: + if self.node_target_type is None: + self.node_target_type = self._metadata[0][0] + x_node = x_dict[self.node_target_type] + batch_node = batch_dict[self.node_target_type] + + try: + head_device = next(headloc.parameters()).device + except StopIteration: + head_device = x_node.device + if x_node.device != head_device: + x_node = x_node.to(head_device) + if batch_node.device != head_device: + batch_node = batch_node.to(head_device) + + node_NN_type = self.config_heads["node"][0]["architecture"]["type"] + if node_NN_type not in ("mlp", "mlp_per_node", "conv"): + raise ValueError( + "HeteroBase only supports node heads with 'mlp', 'mlp_per_node', or 'conv'." + ) + + head = torch.zeros((x_node.shape[0], head_dim), device=x_node.device) + headvar = torch.zeros( + (x_node.shape[0], head_dim * self.var_output), device=x_node.device + ) + + if node_NN_type == "conv": + if self.num_branches != 1: + raise NotImplementedError( + "conv-based node heads with multiple branches are not supported yet for hetero models." + ) + branchtype = "branch-0" + x_dict_node = x_dict + for conv, batch_norm in zip( + headloc[branchtype][0::2], headloc[branchtype][1::2] + ): + if edge_attr_dict is None: + x_dict_node = conv(x_dict_node, data.edge_index_dict) + else: + x_dict_node = conv( + x_dict_node, data.edge_index_dict, edge_attr_dict + ) + for node_type, x in x_dict_node.items(): + x = batch_norm[node_type](x) + x = self.activation_function(x) + x_dict_node[node_type] = x + x_node_out = x_dict_node[self.node_target_type] + head = x_node_out[:, :head_dim] + headvar = x_node_out[:, head_dim:] ** 2 + else: + if self.num_branches == 1: + branchtype = "branch-0" + x_node_out = headloc[branchtype](x=x_node, batch=batch_node) + head = x_node_out[:, :head_dim] + headvar = x_node_out[:, head_dim:] ** 2 + else: + unique, node_counts = torch.unique_consecutive( + batch_node, return_counts=True + ) + for ID in datasetIDs: + mask = data.dataset_name == ID + mask = mask[:, 0] + branchtype = f"branch-{ID.item()}" + node_count = node_counts[unique == ID] + if node_count.numel() == 0: + continue + idx_start = torch.sum(node_counts[unique < ID]) + idx_end = idx_start + node_count[0] + x_node_out = headloc[branchtype]( + x=x_node[idx_start:idx_end], + batch=batch_node[idx_start:idx_end], + ) + head[idx_start:idx_end] = x_node_out[:, :head_dim] + headvar[idx_start:idx_end] = x_node_out[:, head_dim:] ** 2 + + outputs.append(head) + outputs_var.append(headvar) + + return outputs + + def loss(self, pred, value, head_index): + var = None + if self.var_output: + var = pred[1] + pred = pred[0] + if self.ilossweights_nll == 1: + raise ValueError("loss_nll() not ready yet") + if self.ilossweights_hyperp == 1: + return self.loss_hpweighted(pred, value, head_index, var=var) + raise ValueError("Unsupported loss weighting configuration") + + def loss_hpweighted(self, pred, value, head_index, var=None): + tot_loss = 0 + tasks_loss = [] + for ihead in range(self.num_heads): + head_pre = pred[ihead] + pred_shape = head_pre.shape + head_val = value[head_index[ihead]] + value_shape = head_val.shape + if pred_shape != value_shape: + head_val = torch.reshape(head_val, pred_shape) + head_val = head_val.to(head_pre.device) + if var is None: + tot_loss += ( + self.loss_function(head_pre, head_val) * self.loss_weights[ihead] + ) + tasks_loss.append(self.loss_function(head_pre, head_val)) + else: + head_var = var[ihead] + tot_loss += ( + self.loss_function(head_pre, head_val, head_var) + * self.loss_weights[ihead] + ) + tasks_loss.append(self.loss_function(head_pre, head_val, head_var)) + + return tot_loss, tasks_loss diff --git a/hydragnn/models/heterogeneous/HeteroGATStack.py b/hydragnn/models/heterogeneous/HeteroGATStack.py new file mode 100644 index 000000000..aed877f12 --- /dev/null +++ b/hydragnn/models/heterogeneous/HeteroGATStack.py @@ -0,0 +1,183 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +from torch.nn import ModuleDict, ModuleList +from torch_geometric.nn import BatchNorm, GATv2Conv, HeteroConv + +from .HeteroBase import HeteroBase + + +class HeteroGATStack(HeteroBase): + def __init__( + self, + heads: int, + negative_slope: float, + edge_dim: int, + *args, + **kwargs, + ): + self.heads = heads + self.negative_slope = negative_slope + self.edge_dim = edge_dim + self.is_edge_model = True + super().__init__(*args, **kwargs) + + def _build_hetero_gat_conv(self, input_dim: int, output_dim: int, concat: bool): + conv_dict = {} + shared_conv = None + for edge_type in self._metadata[1]: + resolved_edge_dim = self._resolve_edge_dim_for_type(edge_type) + if self.share_relation_weights: + if shared_conv is None: + shared_conv = GATv2Conv( + in_channels=input_dim, + out_channels=output_dim, + heads=self.heads, + negative_slope=self.negative_slope, + dropout=self.dropout, + add_self_loops=False, + edge_dim=resolved_edge_dim, + concat=concat, + ) + conv_dict[edge_type] = shared_conv + else: + conv_dict[edge_type] = GATv2Conv( + in_channels=input_dim, + out_channels=output_dim, + heads=self.heads, + negative_slope=self.negative_slope, + dropout=self.dropout, + add_self_loops=False, + edge_dim=resolved_edge_dim, + concat=concat, + ) + return HeteroConv(conv_dict, aggr="sum") + + def _init_conv(self): + self.graph_convs = ModuleList() + self.feature_layers = ModuleList() + + # First layer: concat=True -> hidden_dim * heads + self.graph_convs.append( + self._build_hetero_gat_conv(self.hidden_dim, self.hidden_dim, True) + ) + node_norms = ModuleDict({}) + for node_type in self._metadata[0]: + node_norms[node_type] = BatchNorm(self.hidden_dim * self.heads) + self.feature_layers.append(node_norms) + + # Middle layers: concat=True -> hidden_dim * heads + for _ in range(self.num_conv_layers - 2): + self.graph_convs.append( + self._build_hetero_gat_conv( + self.hidden_dim * self.heads, self.hidden_dim, True + ) + ) + node_norms = ModuleDict({}) + for node_type in self._metadata[0]: + node_norms[node_type] = BatchNorm(self.hidden_dim * self.heads) + self.feature_layers.append(node_norms) + + # Final layer: concat=False -> hidden_dim + if self.num_conv_layers > 1: + self.graph_convs.append( + self._build_hetero_gat_conv( + self.hidden_dim * self.heads, self.hidden_dim, False + ) + ) + node_norms = ModuleDict({}) + for node_type in self._metadata[0]: + node_norms[node_type] = BatchNorm(self.hidden_dim) + self.feature_layers.append(node_norms) + + self._initialized = True + if "node" not in self.config_heads: + return + nodeconfiglist = self.config_heads["node"] + assert ( + self.num_branches == len(nodeconfiglist) or self.num_branches == 1 + ), "assuming node head has the same branches as graph head, if any" + for branchdict in nodeconfiglist: + if branchdict["architecture"]["type"] != "conv": + return + + node_feature_ind = [ + i for i, head_type in enumerate(self.head_type) if head_type == "node" + ] + if len(node_feature_ind) == 0: + return + + for branchdict in nodeconfiglist: + branchtype = branchdict["type"] + brancharct = branchdict["architecture"] + num_conv_layers_node = brancharct["num_headlayers"] + hidden_dim_node = brancharct["dim_headlayers"] + + convs_node_hidden = ModuleList() + batch_norms_node_hidden = ModuleList() + convs_node_output = ModuleList() + batch_norms_node_output = ModuleList() + + convs_node_hidden.append( + self._build_hetero_gat_conv(self.hidden_dim, hidden_dim_node[0], True) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm(hidden_dim_node[0] * self.heads) + batch_norms_node_hidden.append(bn_dict) + + for ilayer in range(num_conv_layers_node - 1): + convs_node_hidden.append( + self._build_hetero_gat_conv( + hidden_dim_node[ilayer] * self.heads, + hidden_dim_node[ilayer + 1], + True, + ) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm( + hidden_dim_node[ilayer + 1] * self.heads + ) + batch_norms_node_hidden.append(bn_dict) + + for ihead in node_feature_ind: + convs_node_output.append( + self._build_hetero_gat_conv( + hidden_dim_node[-1] * self.heads, + self.head_dims[ihead], + False, + ) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm(self.head_dims[ihead]) + batch_norms_node_output.append(bn_dict) + + self.convs_node_hidden[branchtype] = convs_node_hidden + self.batch_norms_node_hidden[branchtype] = batch_norms_node_hidden + self.convs_node_output[branchtype] = convs_node_output + self.batch_norms_node_output[branchtype] = batch_norms_node_output + + def get_conv(self, input_dim, output_dim, edge_dim=None): + return GATv2Conv( + in_channels=input_dim, + out_channels=output_dim, + heads=self.heads, + negative_slope=self.negative_slope, + dropout=self.dropout, + add_self_loops=False, + edge_dim=edge_dim, + concat=False, + ) + + def __str__(self): + return "HeteroGATStack" diff --git a/hydragnn/models/heterogeneous/HeteroGINStack.py b/hydragnn/models/heterogeneous/HeteroGINStack.py new file mode 100644 index 000000000..6b08e5929 --- /dev/null +++ b/hydragnn/models/heterogeneous/HeteroGINStack.py @@ -0,0 +1,36 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +import torch.nn as nn +from torch_geometric.nn import GINConv + +from .HeteroBase import HeteroBase + + +class HeteroGINStack(HeteroBase): + def __init__(self, *args, **kwargs): + self.is_edge_model = False + super().__init__(*args, **kwargs) + + def get_conv(self, input_dim, output_dim, edge_dim=None): + gin = GINConv( + nn.Sequential( + nn.Linear(input_dim, output_dim), + nn.ReLU(), + nn.Linear(output_dim, output_dim), + ), + eps=100.0, + train_eps=True, + ) + return gin + + def __str__(self): + return "HeteroGINStack" diff --git a/hydragnn/models/heterogeneous/HeteroHEATStack.py b/hydragnn/models/heterogeneous/HeteroHEATStack.py new file mode 100644 index 000000000..7289907f6 --- /dev/null +++ b/hydragnn/models/heterogeneous/HeteroHEATStack.py @@ -0,0 +1,171 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +import torch +from torch.nn import Linear, ModuleDict, ModuleList +from torch_geometric.data import HeteroData +from torch_geometric.nn import BatchNorm, HEATConv + +from .HeteroBase import HeteroBase + + +class HeteroHEATStack(HeteroBase): + def __init__( + self, + attention_heads: int, + edge_type_emb_dim: int, + edge_attr_emb_dim: int, + *args, + **kwargs, + ): + self.attention_heads = attention_heads + self.edge_type_emb_dim = edge_type_emb_dim + self.edge_attr_emb_dim = edge_attr_emb_dim + self.node_types = None + self.edge_types = None + self._heat_edge_dim = None + self.is_edge_model = True + super().__init__(*args, **kwargs) + self.edge_lin_dict = ModuleDict() + + def _init_conv(self): + self.graph_convs = ModuleList() + self.feature_layers = ModuleList() + + self.node_types = list(self._metadata[0]) + self.edge_types = list(self._metadata[1]) + self._heat_edge_dim = self.hidden_dim + + for _ in range(self.num_conv_layers): + self.graph_convs.append( + HEATConv( + in_channels=self.hidden_dim, + out_channels=self.hidden_dim, + num_node_types=len(self.node_types), + num_edge_types=len(self.edge_types), + edge_type_emb_dim=self.edge_type_emb_dim, + edge_dim=self._heat_edge_dim, + edge_attr_emb_dim=self.edge_attr_emb_dim, + heads=self.attention_heads, + concat=False, + ) + ) + node_norms = ModuleDict({}) + for node_type in self.node_types: + node_norms[node_type] = BatchNorm(self.hidden_dim) + self.feature_layers.append(node_norms) + self._initialized = True + + def _init_node_conv(self): + nodeconfiglist = self.config_heads.get("node", []) + for branchdict in nodeconfiglist: + if branchdict["architecture"]["type"] == "conv": + raise NotImplementedError( + "HeteroHEATStack does not support conv-based node heads. Use 'mlp' or 'mlp_per_node'." + ) + + def _ensure_edge_projector(self, edge_type, edge_attr_dim: int, device): + key = str(edge_type) + if key not in self.edge_lin_dict: + self.edge_lin_dict[key] = Linear(edge_attr_dim, self.hidden_dim) + if self.edge_lin_dict[key].weight.device != device: + self.edge_lin_dict[key] = self.edge_lin_dict[key].to(device) + + def forward(self, data): + self._maybe_init_metadata(data) + + x_dict = data.x_dict + self._ensure_node_embedders(x_dict) + x_dict = { + node_type: self.node_embedders[node_type](x.float()) + for node_type, x in x_dict.items() + } + + batch_dict = self._get_batch_dict(data, x_dict) + edge_attr_dict = self._get_edge_attr_dict(data) + + node_heads = self.config_heads.get("node", []) + if node_heads and node_heads[0]["architecture"]["type"] == "conv": + raise NotImplementedError( + "HeteroHEATStack does not support conv-based node heads. Use 'mlp' or 'mlp_per_node'." + ) + + projected_edge_attr_dict = {} + if edge_attr_dict is not None: + for edge_type, edge_attr in edge_attr_dict.items(): + self._ensure_edge_projector( + edge_type, edge_attr.size(-1), edge_attr.device + ) + projected_edge_attr_dict[edge_type] = self.activation_function( + self.edge_lin_dict[str(edge_type)](edge_attr) + ) + + # Build temporary HeteroData for HEATConv -> homogeneous conversion + hdata = HeteroData() + for node_type in self.node_types: + hdata[node_type].x = x_dict[node_type] + + for edge_type in self.edge_types: + if edge_type in data.edge_index_dict: + ei = data.edge_index_dict[edge_type] + else: + device = x_dict[self.node_types[0]].device + ei = torch.empty((2, 0), dtype=torch.long, device=device) + + hdata[edge_type].edge_index = ei + num_edges = ei.size(1) + device = ei.device + + if edge_type in projected_edge_attr_dict: + ea = projected_edge_attr_dict[edge_type] + if ea.size(0) != num_edges: + raise ValueError( + f"edge_attr rows ({ea.size(0)}) must match num_edges ({num_edges}) for edge_type={edge_type}" + ) + hdata[edge_type].edge_attr = ea + else: + hdata[edge_type].edge_attr = torch.zeros( + (num_edges, self._heat_edge_dim), + dtype=hdata[self.node_types[0]].x.dtype, + device=device, + ) + + homo = hdata.to_homogeneous(node_attrs=["x"], edge_attrs=["edge_attr"]) + x = homo.x + edge_index = homo.edge_index + node_type = homo.node_type + edge_type = homo.edge_type + edge_attr = homo.edge_attr + + for conv, node_norms in zip(self.graph_convs, self.feature_layers): + x = conv(x, edge_index, node_type, edge_type, edge_attr) + for idx, node_name in enumerate(self.node_types): + mask = node_type == idx + if not torch.any(mask): + continue + x_type = x[mask] + x_type = self._apply_graph_conditioning( + x_type, batch_dict[node_name], data + ) + x_type = node_norms[node_name](x_type) + x_type = self.activation_function(x_type) + x[mask] = x_type + + # Reconstruct x_dict from homogeneous representation + x_dict = {} + for idx, node_name in enumerate(self.node_types): + mask = node_type == idx + x_dict[node_name] = x[mask] + + return self._decode_from_x_dict(x_dict, batch_dict, data, edge_attr_dict=None) + + def __str__(self): + return "HeteroHEATStack" diff --git a/hydragnn/models/heterogeneous/HeteroHGTStack.py b/hydragnn/models/heterogeneous/HeteroHGTStack.py new file mode 100644 index 000000000..23445ccc6 --- /dev/null +++ b/hydragnn/models/heterogeneous/HeteroHGTStack.py @@ -0,0 +1,77 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +from torch.nn import ModuleDict, ModuleList +from torch_geometric.nn import BatchNorm, HGTConv + +from .HeteroBase import HeteroBase + + +class HeteroHGTStack(HeteroBase): + def __init__(self, num_heads: int, *args, **kwargs): + self.num_heads = num_heads + self.is_edge_model = False + super().__init__(*args, **kwargs) + + def _init_conv(self): + self.graph_convs = ModuleList() + self.feature_layers = ModuleList() + + for _ in range(self.num_conv_layers): + self.graph_convs.append( + HGTConv( + self.hidden_dim, self.hidden_dim, self._metadata, self.num_heads + ) + ) + node_norms = ModuleDict({}) + for node_type in self._metadata[0]: + node_norms[node_type] = BatchNorm(self.hidden_dim) + self.feature_layers.append(node_norms) + self._initialized = True + + def _init_node_conv(self): + nodeconfiglist = self.config_heads.get("node", []) + for branchdict in nodeconfiglist: + if branchdict["architecture"]["type"] == "conv": + raise NotImplementedError( + "HeteroHGTStack does not support conv-based node heads. Use 'mlp' or 'mlp_per_node'." + ) + + def forward(self, data): + self._maybe_init_metadata(data) + + x_dict = data.x_dict + self._ensure_node_embedders(x_dict) + x_dict = { + node_type: self.node_embedders[node_type](x.float()) + for node_type, x in x_dict.items() + } + + batch_dict = self._get_batch_dict(data, x_dict) + + node_heads = self.config_heads.get("node", []) + if node_heads and node_heads[0]["architecture"]["type"] == "conv": + raise NotImplementedError( + "HeteroHGTStack does not support conv-based node heads. Use 'mlp' or 'mlp_per_node'." + ) + + for conv, node_norms in zip(self.graph_convs, self.feature_layers): + x_dict = conv(x_dict, data.edge_index_dict) + for node_type, x in x_dict.items(): + x = self._apply_graph_conditioning(x, batch_dict[node_type], data) + x = node_norms[node_type](x) + x = self.activation_function(x) + x_dict[node_type] = x + + return self._decode_from_x_dict(x_dict, batch_dict, data, edge_attr_dict=None) + + def __str__(self): + return "HeteroHGTStack" diff --git a/hydragnn/models/heterogeneous/HeteroPNAStack.py b/hydragnn/models/heterogeneous/HeteroPNAStack.py new file mode 100644 index 000000000..0f72f75e4 --- /dev/null +++ b/hydragnn/models/heterogeneous/HeteroPNAStack.py @@ -0,0 +1,74 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +import torch +from torch_geometric.nn import PNAConv +from torch.nn import Module + +from .HeteroBase import HeteroBase + + +class _BipartitePNAWrapper(Module): + def __init__(self, conv: PNAConv): + super().__init__() + self.conv = conv + + def forward(self, x, edge_index, edge_attr=None): + if isinstance(x, tuple): + x_src, x_dst = x + x_cat = torch.cat([x_src, x_dst], dim=0) + edge_index_cat = torch.stack( + [edge_index[0], edge_index[1] + x_src.size(0)], dim=0 + ) + out = self.conv(x_cat, edge_index_cat, edge_attr=edge_attr) + return out[x_src.size(0) :] + return self.conv(x, edge_index, edge_attr=edge_attr) + + +class HeteroPNAStack(HeteroBase): + def __init__( + self, + deg: list, + edge_dim: int, + *args, + **kwargs, + ): + self.aggregators = ["mean", "min", "max", "std"] + self.scalers = [ + "identity", + "amplification", + "attenuation", + "linear", + ] + self.deg = torch.Tensor(deg) + self.edge_dim = edge_dim + self.is_edge_model = True + super().__init__(*args, **kwargs) + + def get_conv(self, input_dim, output_dim, edge_dim=None): + # Use the per-type edge_dim passed by _build_hetero_conv. + # For uniform (int) edge_dim, _resolve_edge_dim_for_type returns the + # same int for every edge type, so this is backward-compatible. + conv = PNAConv( + in_channels=input_dim, + out_channels=output_dim, + aggregators=self.aggregators, + scalers=self.scalers, + deg=self.deg, + edge_dim=edge_dim, + pre_layers=1, + post_layers=1, + divide_input=False, + ) + return _BipartitePNAWrapper(conv) + + def __str__(self): + return "HeteroPNAStack" diff --git a/hydragnn/models/heterogeneous/HeteroRGATStack.py b/hydragnn/models/heterogeneous/HeteroRGATStack.py new file mode 100644 index 000000000..f0ce490d4 --- /dev/null +++ b/hydragnn/models/heterogeneous/HeteroRGATStack.py @@ -0,0 +1,171 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +from torch.nn import ModuleDict, ModuleList +from torch_geometric.nn import BatchNorm, GATConv, HeteroConv + +from .HeteroBase import HeteroBase + + +class HeteroRGATStack(HeteroBase): + def __init__( + self, + heads: int, + negative_slope: float, + edge_dim: int, + *args, + **kwargs, + ): + self.heads = heads + self.negative_slope = negative_slope + self.edge_dim = edge_dim + self.is_edge_model = True + super().__init__(*args, **kwargs) + + def _build_hetero_rgat_conv(self, input_dim: int, output_dim: int, concat: bool): + conv_dict = {} + shared_conv = None + for edge_type in self._metadata[1]: + edge_dim = self._resolve_edge_dim_for_type(edge_type) + if self.share_relation_weights: + if shared_conv is None: + shared_conv = GATConv( + in_channels=input_dim, + out_channels=output_dim, + heads=self.heads, + negative_slope=self.negative_slope, + dropout=self.dropout, + add_self_loops=False, + edge_dim=edge_dim, + concat=concat, + ) + conv_dict[edge_type] = shared_conv + else: + conv_dict[edge_type] = GATConv( + in_channels=input_dim, + out_channels=output_dim, + heads=self.heads, + negative_slope=self.negative_slope, + dropout=self.dropout, + add_self_loops=False, + edge_dim=edge_dim, + concat=concat, + ) + return HeteroConv(conv_dict, aggr="sum") + + def _init_conv(self): + self.graph_convs = ModuleList() + self.feature_layers = ModuleList() + + # First layer: concat=True -> hidden_dim * heads + self.graph_convs.append( + self._build_hetero_rgat_conv(self.hidden_dim, self.hidden_dim, True) + ) + node_norms = ModuleDict({}) + for node_type in self._metadata[0]: + node_norms[node_type] = BatchNorm(self.hidden_dim * self.heads) + self.feature_layers.append(node_norms) + + # Middle layers: concat=True -> hidden_dim * heads + for _ in range(self.num_conv_layers - 2): + self.graph_convs.append( + self._build_hetero_rgat_conv( + self.hidden_dim * self.heads, self.hidden_dim, True + ) + ) + node_norms = ModuleDict({}) + for node_type in self._metadata[0]: + node_norms[node_type] = BatchNorm(self.hidden_dim * self.heads) + self.feature_layers.append(node_norms) + + # Final layer: concat=False -> hidden_dim + if self.num_conv_layers > 1: + self.graph_convs.append( + self._build_hetero_rgat_conv( + self.hidden_dim * self.heads, self.hidden_dim, False + ) + ) + node_norms = ModuleDict({}) + for node_type in self._metadata[0]: + node_norms[node_type] = BatchNorm(self.hidden_dim) + self.feature_layers.append(node_norms) + + self._initialized = True + + def _init_node_conv(self): + nodeconfiglist = self.config_heads["node"] + assert ( + self.num_branches == len(nodeconfiglist) or self.num_branches == 1 + ), "assuming node head has the same branches as graph head, if any" + for branchdict in nodeconfiglist: + if branchdict["architecture"]["type"] != "conv": + return + + node_feature_ind = [ + i for i, head_type in enumerate(self.head_type) if head_type == "node" + ] + if len(node_feature_ind) == 0: + return + + for branchdict in nodeconfiglist: + branchtype = branchdict["type"] + brancharct = branchdict["architecture"] + num_conv_layers_node = brancharct["num_headlayers"] + hidden_dim_node = brancharct["dim_headlayers"] + + convs_node_hidden = ModuleList() + batch_norms_node_hidden = ModuleList() + convs_node_output = ModuleList() + batch_norms_node_output = ModuleList() + + convs_node_hidden.append( + self._build_hetero_rgat_conv(self.hidden_dim, hidden_dim_node[0], True) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm(hidden_dim_node[0] * self.heads) + batch_norms_node_hidden.append(bn_dict) + + for ilayer in range(num_conv_layers_node - 1): + convs_node_hidden.append( + self._build_hetero_rgat_conv( + hidden_dim_node[ilayer] * self.heads, + hidden_dim_node[ilayer + 1], + True, + ) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm( + hidden_dim_node[ilayer + 1] * self.heads + ) + batch_norms_node_hidden.append(bn_dict) + + for ihead in node_feature_ind: + convs_node_output.append( + self._build_hetero_rgat_conv( + hidden_dim_node[-1] * self.heads, + self.head_dims[ihead], + False, + ) + ) + bn_dict = ModuleDict({}) + for node_type in self._metadata[0]: + bn_dict[node_type] = BatchNorm(self.head_dims[ihead]) + batch_norms_node_output.append(bn_dict) + + self.convs_node_hidden[branchtype] = convs_node_hidden + self.batch_norms_node_hidden[branchtype] = batch_norms_node_hidden + self.convs_node_output[branchtype] = convs_node_output + self.batch_norms_node_output[branchtype] = batch_norms_node_output + + def __str__(self): + return "HeteroRGATStack" diff --git a/hydragnn/models/heterogeneous/HeteroSAGEStack.py b/hydragnn/models/heterogeneous/HeteroSAGEStack.py new file mode 100644 index 000000000..966876f53 --- /dev/null +++ b/hydragnn/models/heterogeneous/HeteroSAGEStack.py @@ -0,0 +1,26 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +from torch_geometric.nn import SAGEConv + +from .HeteroBase import HeteroBase + + +class HeteroSAGEStack(HeteroBase): + def __init__(self, *args, **kwargs): + self.is_edge_model = False + super().__init__(*args, **kwargs) + + def get_conv(self, input_dim, output_dim, edge_dim=None): + return SAGEConv(in_channels=input_dim, out_channels=output_dim) + + def __str__(self): + return "HeteroSAGEStack" diff --git a/hydragnn/models/heterogeneous/__init__.py b/hydragnn/models/heterogeneous/__init__.py new file mode 100644 index 000000000..d3b14bfaf --- /dev/null +++ b/hydragnn/models/heterogeneous/__init__.py @@ -0,0 +1,8 @@ +from .HeteroBase import HeteroBase +from .HeteroGINStack import HeteroGINStack +from .HeteroSAGEStack import HeteroSAGEStack +from .HeteroGATStack import HeteroGATStack +from .HeteroPNAStack import HeteroPNAStack +from .HeteroHGTStack import HeteroHGTStack +from .HeteroHEATStack import HeteroHEATStack +from .HeteroRGATStack import HeteroRGATStack diff --git a/hydragnn/postprocess/visualizer.py b/hydragnn/postprocess/visualizer.py index 83ee4701f..b4849bc7a 100644 --- a/hydragnn/postprocess/visualizer.py +++ b/hydragnn/postprocess/visualizer.py @@ -103,6 +103,30 @@ def __err_condmean(self, data1, data2, weight=1.0): ) return xcen_pasr, mean1d_cond + def __point_density(self, x, y, bins=60): + x = np.asarray(x).ravel() + y = np.asarray(y).ravel() + + mask = np.isfinite(x) & np.isfinite(y) + density = np.zeros_like(x, dtype=float) + if not np.any(mask): + return density + + x_valid = x[mask] + y_valid = y[mask] + hist2d, xedges, yedges = np.histogram2d(x_valid, y_valid, bins=bins) + + xbin = np.clip(np.digitize(x_valid, xedges) - 1, 0, hist2d.shape[0] - 1) + ybin = np.clip(np.digitize(y_valid, yedges) - 1, 0, hist2d.shape[1] - 1) + density_valid = hist2d[xbin, ybin] + + max_density = np.max(density_valid) + if max_density > 0: + density_valid = density_valid / max_density + + density[mask] = density_valid + return density + def __scatter_impl( self, ax, @@ -116,8 +140,26 @@ def __scatter_impl( y_label=None, xylim_equal=False, ): + x = np.asarray(x).ravel() + y = np.asarray(y).ravel() + marker = "o" if marker is None else marker + point_size = 18 if s is None else s + + density = self.__point_density(x, y) + plot_order = np.argsort(density) + scatter = ax.scatter( + x[plot_order], + y[plot_order], + s=point_size, + c=density[plot_order], + cmap="viridis", + marker=marker, + edgecolors="none", + alpha=0.85, + ) - ax.scatter(x, y, s=s, edgecolor="b", marker=marker, facecolor="none") + colorbar = ax.figure.colorbar(scatter, ax=ax, fraction=0.046, pad=0.04) + colorbar.set_label("Point concentration") ax.set_title(title + ", number of samples =" + str(len(x))) ax.set_xlabel(x_label) diff --git a/hydragnn/train/train_validate_test.py b/hydragnn/train/train_validate_test.py index 035d0cab2..16c93d074 100644 --- a/hydragnn/train/train_validate_test.py +++ b/hydragnn/train/train_validate_test.py @@ -75,13 +75,26 @@ def move_batch_to_device(data, param_dtype): device = get_device() if isinstance(data, torch.Tensor): - data = data.to(dtype=param_dtype) + return data.to(device=device, dtype=param_dtype) else: - for key, value in data.items(): - if isinstance(value, torch.Tensor) and torch.is_floating_point(value): - data[key] = value.to(dtype=param_dtype) - - return data.to(device) + data = data.to(device) + if hasattr(data, "stores"): + for store in data.stores: + for key, value in store.items(): + if isinstance(value, torch.Tensor): + value = value.to(device=device) + if torch.is_floating_point(value): + value = value.to(dtype=param_dtype) + store[key] = value + else: + for key, value in data.items(): + if isinstance(value, torch.Tensor): + value = value.to(device=device) + if torch.is_floating_point(value): + value = value.to(dtype=param_dtype) + data[key] = value + + return data def get_autocast_and_scaler(precision): diff --git a/hydragnn/utils/datasets/__init__.py b/hydragnn/utils/datasets/__init__.py index 6b9beaccf..2416aacc3 100644 --- a/hydragnn/utils/datasets/__init__.py +++ b/hydragnn/utils/datasets/__init__.py @@ -13,6 +13,11 @@ compositional_stratified_splitting, ) from .distdataset import DistDataset + +try: + from .hdf5dataset import HDF5Dataset, HDF5Writer +except ImportError: + pass from .lsmsdataset import LSMSDataset from .pickledataset import SimplePickleDataset, SimplePickleWriter from .serializeddataset import SerializedDataset, SerializedWriter diff --git a/hydragnn/utils/datasets/hdf5dataset.py b/hydragnn/utils/datasets/hdf5dataset.py new file mode 100644 index 000000000..79c1b53ff --- /dev/null +++ b/hydragnn/utils/datasets/hdf5dataset.py @@ -0,0 +1,284 @@ +"""HDF5-based storage for heterogeneous PyG data. + +Each MPI rank writes its own shard file inside a directory: + / + meta.h5 – total counts, rank offsets + shard-0000.h5 – samples from rank 0 + shard-0001.h5 – samples from rank 1 + … + +Each shard stores samples as variable-length byte datasets (one per label) +containing pickle-serialised HeteroData objects. No homogeneous conversion. + +The writer supports two modes of operation: + +**Batch mode** (original API — backward compatible):: + + w = HDF5Writer(basedir, comm) + w.add("trainset", list_of_samples) + w.save() + +**Streaming mode** (memory-efficient — like AdiosWriter):: + + w = HDF5Writer(basedir, comm) + w.begin("trainset") # open a streaming label + for sample in process(...): + w.put(sample) # serialize & flush immediately + w.end_label() # finalize the current label + w.save() # write metadata only +""" + +import os +import pickle + +import h5py +import numpy as np +from mpi4py import MPI + +from .abstractbasedataset import AbstractBaseDataset +from hydragnn.utils.print import iterate_tqdm + + +# ────────────────────────────────────────────────────────────────────── +# Writer +# ────────────────────────────────────────────────────────────────────── +class HDF5Writer: + """Write HDF5 shards with optional streaming to avoid OOM. + + Supports two usage patterns: + + 1. **Batch mode** — ``add(label, data)`` then ``save()`` (original API, + backward compatible). + 2. **Streaming mode** — ``begin(label)`` / ``put(sample)`` / + ``end_label()`` then ``save()``. Samples are serialized and flushed + to disk in small batches; the caller never needs to hold all samples + in memory. + """ + + def __init__(self, basedir, comm=MPI.COMM_WORLD, batch_size=64): + self.basedir = basedir + self.comm = comm + self.rank = comm.Get_rank() + self.size = comm.Get_size() + self.batch_size = batch_size + + # Batch-mode accumulator (used only when add() is called) + self._labels = {} # label -> list[data] + + # Streaming-mode state + self._fh = None # open h5py.File handle + self._shard_path = None + self._stream_label = None # active label being streamed + self._stream_ds = None # active HDF5 dataset + self._stream_buf = [] # pending samples not yet flushed + self._stream_offset = 0 # next write position in the dataset + self._label_counts = {} # label -> final count (for metadata) + self._streaming_used = False + + # ── batch-mode API (backward compatible) ────────────────────────── + + def add(self, label, data): + if label not in self._labels: + self._labels[label] = [] + if isinstance(data, list): + self._labels[label].extend(data) + else: + self._labels[label].append(data) + + # ── streaming-mode API ──────────────────────────────────────────── + + def begin(self, label): + """Start streaming samples for *label*. Opens the shard file on + first call.""" + self._streaming_used = True + if self._stream_label is not None: + self.end_label() + + if self._fh is None: + os.makedirs(self.basedir, exist_ok=True) + self._shard_path = os.path.join(self.basedir, f"shard-{self.rank:04d}.h5") + self._fh = h5py.File(self._shard_path, "w") + + if label in self._fh: + # Resume appending to an existing dataset (e.g. second case + # writing more samples into the same "trainset" dataset). + self._stream_ds = self._fh[label] + self._stream_offset = self._stream_ds.shape[0] + else: + vlen_dt = h5py.vlen_dtype(np.dtype("uint8")) + # Use a resizable (chunked) dataset so we can append without + # knowing the total count up-front. + self._stream_ds = self._fh.create_dataset( + label, + shape=(0,), + maxshape=(None,), + dtype=vlen_dt, + chunks=(self.batch_size,), + ) + self._stream_offset = 0 + self._stream_label = label + self._stream_buf = [] + + def put(self, sample): + """Add a single sample. Automatically flushes to disk every + *batch_size* samples.""" + self._stream_buf.append(sample) + if len(self._stream_buf) >= self.batch_size: + self._flush_stream_buf() + + def end_label(self): + """Finalize the current streaming label — flush remaining samples.""" + if self._stream_buf: + self._flush_stream_buf() + if self._stream_label is not None: + self._label_counts[self._stream_label] = self._stream_offset + self._stream_label = None + self._stream_ds = None + self._stream_buf = [] + + def _flush_stream_buf(self): + """Serialize buffered samples and write them to the open dataset.""" + n = len(self._stream_buf) + if n == 0: + return + new_end = self._stream_offset + n + self._stream_ds.resize((new_end,)) + for j, sample in enumerate(self._stream_buf): + self._stream_ds[self._stream_offset + j] = np.frombuffer( + pickle.dumps(sample, protocol=4), dtype=np.uint8 + ) + self._stream_offset = new_end + self._stream_buf.clear() + + # ── save (works for both modes) ─────────────────────────────────── + + def save(self): + # Finalize any open streaming label + if self._stream_label is not None: + self.end_label() + + # If batch-mode data was collected, write it now + if self._labels: + self._save_batch_mode() + + # Close the shard file if streaming opened it + if self._fh is not None: + self._fh.close() + self._fh = None + + self.comm.Barrier() + + # ── Metadata ───────────────────────────────────────────────── + + # Merge label names from both modes + local_labels = sorted( + set(list(self._labels.keys()) + list(self._label_counts.keys())) + ) + all_labels = self.comm.allgather(local_labels) + labels = sorted({lbl for lst in all_labels for lbl in lst}) + + per_label_counts = {} + for label in labels: + if label in self._label_counts: + local_n = self._label_counts[label] + else: + local_n = len(self._labels.get(label, [])) + counts = self.comm.gather(local_n, root=0) + if self.rank == 0: + per_label_counts[label] = counts + + if self.rank == 0: + meta_path = os.path.join(self.basedir, "meta.h5") + with h5py.File(meta_path, "w") as fh: + fh.attrs["num_shards"] = self.size + for label in labels: + grp = fh.create_group(label) + counts = np.array(per_label_counts[label], dtype=np.int64) + grp.create_dataset("counts", data=counts) + grp.attrs["total"] = int(counts.sum()) + + self.comm.Barrier() + + def _save_batch_mode(self): + """Write all batch-accumulated data (original logic).""" + if self.rank == 0: + os.makedirs(self.basedir, exist_ok=True) + self.comm.Barrier() + + BATCH = self.batch_size + shard_path = os.path.join(self.basedir, f"shard-{self.rank:04d}.h5") + vlen_dt = h5py.vlen_dtype(np.dtype("uint8")) + mode = "a" if self._streaming_used else "w" + with h5py.File(shard_path, mode) as fh: + for label, samples in self._labels.items(): + n = len(samples) + if n > 0: + ds = fh.create_dataset(label, shape=(n,), dtype=vlen_dt) + for start in iterate_tqdm( + range(0, n, BATCH), + 2, + total=(n + BATCH - 1) // BATCH, + desc=f"HDF5 write {label}", + ): + end = min(start + BATCH, n) + for j, sample in enumerate(samples[start:end]): + ds[start + j] = np.frombuffer( + pickle.dumps(sample, protocol=4), + dtype=np.uint8, + ) + # Release processed samples to reduce memory pressure + for k in range(start, end): + samples[k] = None + else: + fh.create_dataset(label, shape=(0,), dtype=vlen_dt) + + +# ────────────────────────────────────────────────────────────────────── +# Dataset (reader) +# ────────────────────────────────────────────────────────────────────── +class HDF5Dataset(AbstractBaseDataset): + """Read back a split written by HDF5Writer.""" + + def __init__(self, basedir, label, var_config=None): + super().__init__() + self.basedir = basedir + self.label = label + self.var_config = var_config + + meta_path = os.path.join(basedir, "meta.h5") + with h5py.File(meta_path, "r") as fh: + grp = fh[label] + self.counts = grp["counts"][:] + self.ntotal = int(grp.attrs["total"]) + + # Build a mapping from global index -> (shard_rank, local_index) + self._offsets = np.zeros(len(self.counts) + 1, dtype=np.int64) + np.cumsum(self.counts, out=self._offsets[1:]) + + # Cache open file handles lazily + self._handles = {} + + def len(self): + return self.ntotal + + def get(self, idx): + # Find which shard this global index belongs to + shard = int(np.searchsorted(self._offsets[1:], idx, side="right")) + local_idx = idx - int(self._offsets[shard]) + fh = self._open_shard(shard) + blob = fh[self.label][local_idx] + data = pickle.loads(blob.tobytes()) + return data + + def _open_shard(self, shard): + if shard not in self._handles: + path = os.path.join(self.basedir, f"shard-{shard:04d}.h5") + self._handles[shard] = h5py.File(path, "r") + return self._handles[shard] + + def __del__(self): + for fh in self._handles.values(): + try: + fh.close() + except Exception: + pass diff --git a/hydragnn/utils/distributed/distributed.py b/hydragnn/utils/distributed/distributed.py index b348ef014..064cff436 100644 --- a/hydragnn/utils/distributed/distributed.py +++ b/hydragnn/utils/distributed/distributed.py @@ -76,6 +76,13 @@ def find_ifname(myaddr): return ifname +def _resolve_ipv4_addr(host: str) -> str: + try: + return socket.gethostbyname(host) + except socket.gaierror: + return host + + def parse_slurm_nodelist(nodelist): """ Parse SLURM_NODELIST env string to get list of nodes. @@ -214,6 +221,9 @@ def _derive_master_port(default_port: int = 8889) -> str: ## The following is CADES specific master_addr = parse_slurm_nodelist(os.environ["PBS_O_HOST"])[0] + if os.getenv("HYDRAGNN_PREFER_IPV4", "1") == "1": + master_addr = _resolve_ipv4_addr(master_addr) + try: port_retries = int(os.getenv("HYDRAGNN_MASTER_PORT_RETRIES", "8")) explicit_port = os.getenv("HYDRAGNN_MASTER_PORT") is not None diff --git a/hydragnn/utils/input_config_parsing/config_utils.py b/hydragnn/utils/input_config_parsing/config_utils.py index be71925c7..9d82413e9 100644 --- a/hydragnn/utils/input_config_parsing/config_utils.py +++ b/hydragnn/utils/input_config_parsing/config_utils.py @@ -37,6 +37,22 @@ def update_config(config, train_loader, val_loader, test_loader): if "Dataset" in config: check_output_dim_consistent(train_loader.dataset[0], config) + # Always sync node_input_dims from heterogeneous data. + arch_cfg = config["NeuralNetwork"].setdefault("Architecture", {}) + data_sample = train_loader.dataset[0] + if hasattr(data_sample, "node_types"): + node_input_dims = {} + for node_type in data_sample.node_types: + node_store = data_sample[node_type] + if hasattr(node_store, "x") and node_store.x is not None: + node_input_dims[str(node_type)] = int(node_store.x.shape[1]) + if node_input_dims: + if arch_cfg.get("node_input_dims") not in (None, node_input_dims): + warnings.warn( + "Overriding node_input_dims with dataset-derived sizes for hetero model." + ) + arch_cfg["node_input_dims"] = node_input_dims + # Set default values for GPS variables if "global_attn_engine" not in config["NeuralNetwork"]["Architecture"]: config["NeuralNetwork"]["Architecture"]["global_attn_engine"] = None @@ -145,6 +161,15 @@ def update_config(config, train_loader, val_loader, test_loader): if "activation_function" not in config["NeuralNetwork"]["Architecture"]: config["NeuralNetwork"]["Architecture"]["activation_function"] = "relu" + if "hetero_attention_heads" not in config["NeuralNetwork"]["Architecture"]: + config["NeuralNetwork"]["Architecture"]["hetero_attention_heads"] = 4 + if "hetero_attention_negative_slope" not in config["NeuralNetwork"]["Architecture"]: + config["NeuralNetwork"]["Architecture"]["hetero_attention_negative_slope"] = 0.2 + if "hetero_edge_type_emb_dim" not in config["NeuralNetwork"]["Architecture"]: + config["NeuralNetwork"]["Architecture"]["hetero_edge_type_emb_dim"] = 16 + if "hetero_edge_attr_emb_dim" not in config["NeuralNetwork"]["Architecture"]: + config["NeuralNetwork"]["Architecture"]["hetero_edge_attr_emb_dim"] = 16 + if "SyncBatchNorm" not in config["NeuralNetwork"]["Architecture"]: config["NeuralNetwork"]["Architecture"]["SyncBatchNorm"] = False @@ -177,32 +202,64 @@ def update_config_equivariance(config): def update_config_edge_dim(config): - config["edge_dim"] = None - edge_models = [ - "GAT", - "PNA", - "PNAPlus", - "PAINN", - "PNAEq", - "CGCNN", - "SchNet", - "EGNN", - "DimeNet", - "MACE", - ] - if "edge_features" in config and config["edge_features"]: - assert ( - config["mpnn_type"] in edge_models - ), "Edge features can only be used with GAT, PNA, PNAPlus, PAINN, PNAEq, CGCNN, SchNet, EGNN, DimeNet, MACE." - config["edge_dim"] = len(config["edge_features"]) - if "enable_interatomic_potential" in config: - assert not config[ - "enable_interatomic_potential" - ], "Edge features cannot be used with interatomic potentials as the model builds its own specialized features for force computation." - elif config["mpnn_type"] == "CGCNN": - # CG always needs an integer edge_dim - # PNA, PNAPlus, and DimeNet would fail with integer edge_dim without edge_attr - config["edge_dim"] = 0 + def _normalize_edge_dim(value): + if value is None: + return None + if isinstance(value, dict): + # Per-edge-type widths (heterogeneous route). + return {str(k): int(v) for k, v in value.items()} + try: + edge_dim = int(value) + except (TypeError, ValueError) as exc: + raise ValueError( + f"edge_dim must be an integer or dict, got: {value}" + ) from exc + if edge_dim < 0: + raise ValueError(f"edge_dim must be >= 0, got: {edge_dim}") + return edge_dim + + explicit_edge_dim = _normalize_edge_dim(config.get("edge_dim")) + if explicit_edge_dim is not None: + # Explicit edge_dim provided — validate against feature names if any. + if isinstance(explicit_edge_dim, int): + feature_names = config.get("edge_feature_names") + if feature_names: + names_len = len(feature_names) + if names_len != explicit_edge_dim: + raise ValueError( + "NeuralNetwork.Architecture.edge_feature_names length " + f"({names_len}) must match edge_dim ({explicit_edge_dim})." + ) + config["edge_dim"] = explicit_edge_dim + else: + # Backward-compatible auto-computation from edge_features. + edge_models = [ + "GAT", + "PNA", + "PNAPlus", + "PAINN", + "PNAEq", + "CGCNN", + "SchNet", + "EGNN", + "DimeNet", + "MACE", + ] + config["edge_dim"] = None + if "edge_features" in config and config["edge_features"]: + assert ( + config["mpnn_type"] in edge_models + ), "Edge features can only be used with GAT, PNA, PNAPlus, PAINN, PNAEq, CGCNN, SchNet, EGNN, DimeNet, MACE." + config["edge_dim"] = len(config["edge_features"]) + if "enable_interatomic_potential" in config: + assert not config[ + "enable_interatomic_potential" + ], "Edge features cannot be used with interatomic potentials as the model builds its own specialized features for force computation." + elif config["mpnn_type"] == "CGCNN": + # CG always needs an integer edge_dim + # PNA, PNAPlus, and DimeNet would fail with integer edge_dim without edge_attr + config["edge_dim"] = 0 + return config @@ -247,9 +304,14 @@ def update_config_NN_outputs(config, data, graph_size_variable): raise ValueError( '"mlp_per_node" is not allowed for variable graph size, Please set config["NeuralNetwork"]["Architecture"]["output_heads"]["node"]["type"] to be "mlp" or "conv" in input file.' ) + denom_nodes = ( + data.y_num_nodes + if hasattr(data, "y_num_nodes") and data.y_num_nodes is not None + else data.num_nodes + ) dim_item = ( data.y_loc[0, ihead + 1].item() - data.y_loc[0, ihead].item() - ) // data.num_nodes + ) // denom_nodes else: raise ValueError("Unknown output type", output_type[ihead]) dims_list.append(dim_item) diff --git a/hydragnn/utils/model/model.py b/hydragnn/utils/model/model.py index 70c6ac906..e473d15f1 100644 --- a/hydragnn/utils/model/model.py +++ b/hydragnn/utils/model/model.py @@ -57,6 +57,8 @@ def loss_function_selection(loss_function_string: str): return lambda x, y: torch.sqrt(torch.nn.functional.mse_loss(x, y)) elif loss_function_string == "GaussianNLLLoss": return torch.nn.GaussianNLLLoss() + elif loss_function_string == "binary_cross_entropy": + return torch.nn.functional.binary_cross_entropy_with_logits else: ImportError diff --git a/run-scripts/SC25-baseline-singledataset0.sh b/run-scripts/SC25-baseline-singledataset0.sh old mode 100755 new mode 100644 index 369ab6bcb..d3cf689d6 --- a/run-scripts/SC25-baseline-singledataset0.sh +++ b/run-scripts/SC25-baseline-singledataset0.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-baseline-singledataset1.sh b/run-scripts/SC25-baseline-singledataset1.sh old mode 100755 new mode 100644 index debf27a1d..3dcd1b04c --- a/run-scripts/SC25-baseline-singledataset1.sh +++ b/run-scripts/SC25-baseline-singledataset1.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-baseline-singledataset2.sh b/run-scripts/SC25-baseline-singledataset2.sh old mode 100755 new mode 100644 index 18cfc29d9..7285b0e60 --- a/run-scripts/SC25-baseline-singledataset2.sh +++ b/run-scripts/SC25-baseline-singledataset2.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-baseline-singledataset3.sh b/run-scripts/SC25-baseline-singledataset3.sh old mode 100755 new mode 100644 index b2449980d..36b5f3662 --- a/run-scripts/SC25-baseline-singledataset3.sh +++ b/run-scripts/SC25-baseline-singledataset3.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-baseline-singledataset4.sh b/run-scripts/SC25-baseline-singledataset4.sh old mode 100755 new mode 100644 index 360a9a53a..6043a4a3d --- a/run-scripts/SC25-baseline-singledataset4.sh +++ b/run-scripts/SC25-baseline-singledataset4.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-baseline.sh b/run-scripts/SC25-baseline.sh old mode 100755 new mode 100644 index 105a73d69..e19e12a8c --- a/run-scripts/SC25-baseline.sh +++ b/run-scripts/SC25-baseline.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-inference.sh b/run-scripts/SC25-inference.sh old mode 100755 new mode 100644 index c64e85c07..dcceae6f0 --- a/run-scripts/SC25-inference.sh +++ b/run-scripts/SC25-inference.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-job-strong.sh b/run-scripts/SC25-job-strong.sh old mode 100755 new mode 100644 index 93bce7117..fdc475cd7 --- a/run-scripts/SC25-job-strong.sh +++ b/run-scripts/SC25-job-strong.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH echo ">>>>" which python diff --git a/run-scripts/SC25-job-weak.sh b/run-scripts/SC25-job-weak.sh old mode 100755 new mode 100644 index 44ad1db53..3422b040d --- a/run-scripts/SC25-job-weak.sh +++ b/run-scripts/SC25-job-weak.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-multibranch-omnistat.sh b/run-scripts/SC25-multibranch-omnistat.sh old mode 100755 new mode 100644 index 1067b9bfb..13227dc03 --- a/run-scripts/SC25-multibranch-omnistat.sh +++ b/run-scripts/SC25-multibranch-omnistat.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/SC25-multibranch.sh b/run-scripts/SC25-multibranch.sh old mode 100755 new mode 100644 index 284257096..e520b032b --- a/run-scripts/SC25-multibranch.sh +++ b/run-scripts/SC25-multibranch.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-baseline-singledataset0.sh b/run-scripts/job-baseline-singledataset0.sh old mode 100755 new mode 100644 index 376b2dd58..1bb2d5970 --- a/run-scripts/job-baseline-singledataset0.sh +++ b/run-scripts/job-baseline-singledataset0.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-baseline-singledataset1.sh b/run-scripts/job-baseline-singledataset1.sh old mode 100755 new mode 100644 index 6a25c2449..388474e79 --- a/run-scripts/job-baseline-singledataset1.sh +++ b/run-scripts/job-baseline-singledataset1.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-baseline-singledataset2.sh b/run-scripts/job-baseline-singledataset2.sh old mode 100755 new mode 100644 index 41afb7580..887b35fcb --- a/run-scripts/job-baseline-singledataset2.sh +++ b/run-scripts/job-baseline-singledataset2.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-baseline-singledataset3.sh b/run-scripts/job-baseline-singledataset3.sh old mode 100755 new mode 100644 index 7efd3833e..5f59e6792 --- a/run-scripts/job-baseline-singledataset3.sh +++ b/run-scripts/job-baseline-singledataset3.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-baseline-singledataset4.sh b/run-scripts/job-baseline-singledataset4.sh old mode 100755 new mode 100644 index 54650fb21..766e8211f --- a/run-scripts/job-baseline-singledataset4.sh +++ b/run-scripts/job-baseline-singledataset4.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-baseline-singledatasets.sh b/run-scripts/job-baseline-singledatasets.sh old mode 100755 new mode 100644 index 43dff4d2e..6fe72e4d7 --- a/run-scripts/job-baseline-singledatasets.sh +++ b/run-scripts/job-baseline-singledatasets.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-baseline.sh b/run-scripts/job-baseline.sh old mode 100755 new mode 100644 index bf17f5959..fd78de9f2 --- a/run-scripts/job-baseline.sh +++ b/run-scripts/job-baseline.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-multibranch-omnistat.sh b/run-scripts/job-multibranch-omnistat.sh old mode 100755 new mode 100644 index dd0f3e151..573a2f874 --- a/run-scripts/job-multibranch-omnistat.sh +++ b/run-scripts/job-multibranch-omnistat.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/run-scripts/job-multibranch-taskparallel.sh b/run-scripts/job-multibranch-taskparallel.sh old mode 100755 new mode 100644 index 8efde0297..be631a48f --- a/run-scripts/job-multibranch-taskparallel.sh +++ b/run-scripts/job-multibranch-taskparallel.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH ## Score-P module use -a /lustre/orion/world-shared/lrn070/jyc/frontier/sw/modulefiles diff --git a/run-scripts/job-multibranch.sh b/run-scripts/job-multibranch.sh old mode 100755 new mode 100644 index 2a74d8b51..84079f04d --- a/run-scripts/job-multibranch.sh +++ b/run-scripts/job-multibranch.sh @@ -13,10 +13,10 @@ # Load conda environemnt source /lustre/orion/lrn070/world-shared/mlupopa/module-to-load-frontier-rocm624.sh source /lustre/orion/lrn070/world-shared/mlupopa/max_conda_envs_frontier/bin/activate -conda activate hydragnn_rocm624 +source activate /lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv #export python path to use ADIOS2 v.2.9.2 -export PYTHONPATH=/lustre/orion/lrn070/world-shared/mlupopa/ADIOS_ROCm624/adios2-install/lib/python3.11/site-packages/:$PYTHONPATH +export PYTHONPATH=/lustre/orion/lrn078/proj-shared/HydraGNN/installation_DOE_supercomputers/HydraGNN-Installation-Frontier/hydragnn_venv/lib/python3.11/site-packages/:$PYTHONPATH which python python -c "import numpy; print(numpy.__version__)" diff --git a/tests/test_heterogeneous_message_passing.py b/tests/test_heterogeneous_message_passing.py new file mode 100644 index 000000000..f30e6cdb4 --- /dev/null +++ b/tests/test_heterogeneous_message_passing.py @@ -0,0 +1,404 @@ +############################################################################## +# Copyright (c) 2024, Oak Ridge National Laboratory # +# All rights reserved. # +# # +# This file is part of HydraGNN and is distributed under a BSD 3-clause # +# license. For the licensing terms see the LICENSE file in the top-level # +# directory. # +# # +# SPDX-License-Identifier: BSD-3-Clause # +############################################################################## + +import pytest +import torch +from torch_geometric.data import HeteroData +from torch_geometric.loader import DataLoader + +import hydragnn +from hydragnn.models.create import create_model +from hydragnn.utils.model.model import update_multibranch_heads +from hydragnn.preprocess.load_data import split_dataset +from hydragnn.utils.distributed import setup_ddp, get_distributed_model + + +def _build_simple_hetero_graph(input_dim: int = 4, edge_dim: int = None): + data = HeteroData() + + # Node features + data["a"].x = torch.randn(4, input_dim) + data["b"].x = torch.randn(3, input_dim) + + # Edges: a -> b + edge_index_ab = torch.tensor([[0, 1, 2, 3], [0, 1, 1, 2]], dtype=torch.long) + data[("a", "to", "b")].edge_index = edge_index_ab + if edge_dim is not None: + data[("a", "to", "b")].edge_attr = torch.randn(edge_index_ab.size(1), edge_dim) + + # Edges: b -> a + edge_index_ba = torch.tensor([[0, 1, 2], [1, 2, 3]], dtype=torch.long) + data[("b", "to", "a")].edge_index = edge_index_ba + if edge_dim is not None: + data[("b", "to", "a")].edge_attr = torch.randn(edge_index_ba.size(1), edge_dim) + + return data + + +def _build_random_hetero_graph( + rng: torch.Generator, + input_dim: int = 4, + edge_dim: int = None, + num_nodes_a: int = 5, + num_nodes_b: int = 4, +): + data = HeteroData() + + # Graph-level latent signals for higher statistical quality. + z_a = torch.randn(input_dim, generator=rng) + z_b = torch.randn(input_dim, generator=rng) + + data["a"].x = z_a + 0.2 * torch.randn(num_nodes_a, input_dim, generator=rng) + data["b"].x = z_b + 0.2 * torch.randn(num_nodes_b, input_dim, generator=rng) + + num_edges_ab = max(2 * num_nodes_a, 4) + src_ab = torch.randint(0, num_nodes_a, (num_edges_ab,), generator=rng) + dst_ab = torch.randint(0, num_nodes_b, (num_edges_ab,), generator=rng) + src_ab = torch.cat([src_ab, torch.arange(num_nodes_a)]) + dst_ab = torch.cat( + [dst_ab, torch.randint(0, num_nodes_b, (num_nodes_a,), generator=rng)] + ) + data[("a", "to", "b")].edge_index = torch.stack([src_ab, dst_ab], dim=0) + + num_edges_ba = max(2 * num_nodes_b, 4) + src_ba = torch.randint(0, num_nodes_b, (num_edges_ba,), generator=rng) + dst_ba = torch.randint(0, num_nodes_a, (num_edges_ba,), generator=rng) + src_ba = torch.cat([src_ba, torch.arange(num_nodes_b)]) + dst_ba = torch.cat( + [dst_ba, torch.randint(0, num_nodes_a, (num_nodes_b,), generator=rng)] + ) + data[("b", "to", "a")].edge_index = torch.stack([src_ba, dst_ba], dim=0) + + if edge_dim is not None: + edge_index_ab = data[("a", "to", "b")].edge_index + edge_index_ba = data[("b", "to", "a")].edge_index + + xa = data["a"].x[edge_index_ab[0]] + xb = data["b"].x[edge_index_ab[1]] + base_ab = 0.5 * (xa + xb).mean(dim=1, keepdim=True) + data[("a", "to", "b")].edge_attr = base_ab.repeat( + 1, edge_dim + ) + 0.05 * torch.randn(base_ab.size(0), edge_dim, generator=rng) + + xb = data["b"].x[edge_index_ba[0]] + xa = data["a"].x[edge_index_ba[1]] + base_ba = 0.5 * (xa + xb).mean(dim=1, keepdim=True) + data[("b", "to", "a")].edge_attr = base_ba.repeat( + 1, edge_dim + ) + 0.05 * torch.randn(base_ba.size(0), edge_dim, generator=rng) + + edge_index_ab = data[("a", "to", "b")].edge_index + edge_index_ba = data[("b", "to", "a")].edge_index + + edge_signal_ab = ( + data["a"].x[edge_index_ab[0]].mean(dim=1) + * data["b"].x[edge_index_ab[1]].mean(dim=1) + ).mean() + edge_signal_ba = ( + data["b"].x[edge_index_ba[0]].mean(dim=1) + * data["a"].x[edge_index_ba[1]].mean(dim=1) + ).mean() + + edge_attr_signal = 0.0 + if edge_dim is not None: + edge_attr_signal = ( + data[("a", "to", "b")].edge_attr.mean() + + data[("b", "to", "a")].edge_attr.mean() + ) + + graph_value = ( + 0.6 * z_a.mean() + + 0.4 * z_b.mean() + + 0.25 * edge_signal_ab + + 0.25 * edge_signal_ba + + 0.1 * edge_attr_signal + + 0.02 * torch.randn((), generator=rng) + ) + data.y = graph_value.view(1, 1) + + return data + + +def _build_random_hetero_dataset( + num_graphs: int, + input_dim: int, + edge_dim: int, + seed: int = 0, +): + rng = torch.Generator().manual_seed(seed) + dataset = [] + for _ in range(num_graphs): + num_nodes_a = int(torch.randint(3, 7, (1,), generator=rng).item()) + num_nodes_b = int(torch.randint(2, 6, (1,), generator=rng).item()) + dataset.append( + _build_random_hetero_graph( + rng, + input_dim=input_dim, + edge_dim=edge_dim, + num_nodes_a=num_nodes_a, + num_nodes_b=num_nodes_b, + ) + ) + return dataset + + +class _HeteroBatchAdapter: + def __init__(self, loader, node_type="a"): + self.loader = loader + self.node_type = node_type + self.dataset = loader.dataset + self.sampler = loader.sampler + self.batch_size = loader.batch_size + self.drop_last = loader.drop_last + + def __len__(self): + return len(self.loader) + + def __iter__(self): + for data in self.loader: + if not hasattr(data, "batch"): + data.batch = data[self.node_type].batch + yield data + + +@pytest.mark.mpi_skip() +@pytest.mark.parametrize( + "mpnn_type,edge_dim,pna_deg", + [ + ("HeteroGIN", None, None), + ("HeteroSAGE", None, None), + ("HeteroGAT", 3, None), + ("HeteroRGAT", 3, None), + ("HeteroHGT", None, None), + ("HeteroHEAT", 3, None), + ("HeteroPNA", 3, [1, 2, 3, 2]), + ], +) +def pytest_hetero_graph_head_forward(mpnn_type, edge_dim, pna_deg): + data = _build_simple_hetero_graph(edge_dim=edge_dim) + + output_heads = { + "graph": { + "num_sharedlayers": 1, + "dim_sharedlayers": 16, + "num_headlayers": 1, + "dim_headlayers": [8], + } + } + + config_args = { + "mpnn_type": mpnn_type, + "input_dim": 4, + "hidden_dim": 16, + "output_dim": [2], + "pe_dim": 1, + "global_attn_engine": "", + "global_attn_type": "", + "global_attn_heads": 1, + "output_type": ["graph"], + "output_heads": update_multibranch_heads(output_heads), + "activation_function": "relu", + "loss_function_type": "mse", + "task_weights": [1.0], + "num_conv_layers": 2, + "equivariance": False, + "use_graph_attr_conditioning": False, + "graph_pooling": "mean", + "hetero_pooling_mode": "sum", + } + if edge_dim is not None: + config_args["edge_dim"] = edge_dim + if pna_deg is not None: + config_args["pna_deg"] = pna_deg + + model = create_model(**config_args) + model.eval() + + outputs = model(data) + assert isinstance(outputs, list) + assert outputs[0].shape == (1, 2) + + +@pytest.mark.mpi_skip() +@pytest.mark.parametrize( + "mpnn_type,edge_dim,pna_deg", + [ + ("HeteroGIN", None, None), + ("HeteroSAGE", None, None), + ("HeteroGAT", 3, None), + ("HeteroPNA", 3, [1, 2, 3, 2]), + ], +) +def pytest_hetero_node_conv_head_forward(mpnn_type, edge_dim, pna_deg): + data = _build_simple_hetero_graph(edge_dim=edge_dim) + + output_heads = { + "node": { + "num_headlayers": 2, + "dim_headlayers": [16, 8], + "type": "conv", + } + } + + config_args = { + "mpnn_type": mpnn_type, + "input_dim": 4, + "hidden_dim": 16, + "output_dim": [1], + "pe_dim": 1, + "global_attn_engine": "", + "global_attn_type": "", + "global_attn_heads": 1, + "output_type": ["node"], + "output_heads": update_multibranch_heads(output_heads), + "activation_function": "relu", + "loss_function_type": "mse", + "task_weights": [1.0], + "num_conv_layers": 2, + "equivariance": False, + "use_graph_attr_conditioning": False, + "graph_pooling": "mean", + "hetero_pooling_mode": "sum", + "node_target_type": "a", + } + if edge_dim is not None: + config_args["edge_dim"] = edge_dim + if pna_deg is not None: + config_args["pna_deg"] = pna_deg + + model = create_model(**config_args) + model.eval() + + outputs = model(data) + assert isinstance(outputs, list) + assert outputs[0].shape == (data["a"].num_nodes, 1) + + +@pytest.mark.mpi_skip() +@pytest.mark.parametrize( + "mpnn_type,edge_dim,pna_deg", + [ + ("HeteroGIN", None, None), + ("HeteroSAGE", None, None), + ("HeteroGAT", 3, None), + ("HeteroPNA", 3, [1, 2, 3, 2]), + ], +) +def pytest_hetero_mpnn_training_randomized_dataset(mpnn_type, edge_dim, pna_deg): + torch.manual_seed(7) + + dataset = _build_random_hetero_dataset( + num_graphs=30000, + input_dim=4, + edge_dim=edge_dim, + seed=13, + ) + trainset, valset, testset = split_dataset( + dataset, perc_train=0.8, stratify_splitting=False + ) + train_loader, val_loader, test_loader = hydragnn.preprocess.create_dataloaders( + trainset, valset, testset, batch_size=8 + ) + train_loader = _HeteroBatchAdapter(train_loader) + val_loader = _HeteroBatchAdapter(val_loader) + test_loader = _HeteroBatchAdapter(test_loader) + + output_heads = { + "graph": { + "num_sharedlayers": 1, + "dim_sharedlayers": 16, + "num_headlayers": 2, + "dim_headlayers": [16, 8], + } + } + + config_args = { + "mpnn_type": mpnn_type, + "input_dim": 4, + "hidden_dim": 16, + "output_dim": [1], + "pe_dim": 0, + "global_attn_engine": "", + "global_attn_type": "", + "global_attn_heads": 1, + "output_type": ["graph"], + "output_heads": update_multibranch_heads(output_heads), + "activation_function": "relu", + "loss_function_type": "mse", + "task_weights": [1.0], + "num_conv_layers": 2, + "equivariance": False, + "use_graph_attr_conditioning": False, + "graph_pooling": "mean", + "hetero_pooling_mode": "sum", + } + if edge_dim is not None: + config_args["edge_dim"] = edge_dim + if pna_deg is not None: + config_args["pna_deg"] = pna_deg + + verbosity = 0 + setup_ddp() + + model = create_model(**config_args) + model = get_distributed_model(model, verbosity) + + optimizer = torch.optim.AdamW(model.parameters(), lr=0.02) + scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( + optimizer, mode="min", factor=0.5, patience=3, min_lr=1.0e-5 + ) + + nn_config = { + "Training": { + "num_epoch": 10, + "conv_checkpointing": False, + }, + "Variables_of_interest": {"output_names": ["y"]}, + } + + log_name = "hetero_mpnn_randomized" + + hydragnn.train.train_validate_test( + model, + optimizer, + train_loader, + val_loader, + test_loader, + None, + scheduler, + nn_config, + log_name, + verbosity, + create_plots=False, + ) + + final_loss, _, true_values, predicted_values = hydragnn.train.test( + test_loader, + model, + verbosity, + num_tasks=1, + precision="fp32", + ) + + mae = torch.nn.L1Loss() + final_mae = mae(true_values[0], predicted_values[0]) + + thresholds = { + "HeteroGIN": [0.09, 0.26], + "HeteroSAGE": [0.09, 0.27], + "HeteroGAT": [0.13, 0.30], + "HeteroPNA": [0.11, 0.28], + } + + assert torch.isfinite(torch.tensor(final_loss)) + assert torch.isfinite(final_mae) + assert final_loss < thresholds[mpnn_type][0] + assert final_mae < thresholds[mpnn_type][1] diff --git a/tests/test_model_loadpred.py b/tests/test_model_loadpred.py index 02162d9b4..8e7940784 100755 --- a/tests/test_model_loadpred.py +++ b/tests/test_model_loadpred.py @@ -100,6 +100,8 @@ def pytest_model_loadpred(): if not case_exist: unittest_train_model( config["NeuralNetwork"]["Architecture"]["model_type"], + None, + None, "ci_multihead.json", False, False,