From 8d954bd3b05ddaf9c3ce71639801e54273706496 Mon Sep 17 00:00:00 2001 From: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> Date: Fri, 18 Sep 2026 21:15:36 +0100 Subject: [PATCH 1/6] add branch flow prediction Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> --- gridfm_graphkit/cli.py | 2 + gridfm_graphkit/tasks/pf_task.py | 20 ++++++- gridfm_graphkit/tasks/utils.py | 94 ++++++++++++++++++++++++++++++++ 3 files changed, 113 insertions(+), 3 deletions(-) diff --git a/gridfm_graphkit/cli.py b/gridfm_graphkit/cli.py index 9f0a9fac..f8fd0625 100644 --- a/gridfm_graphkit/cli.py +++ b/gridfm_graphkit/cli.py @@ -91,6 +91,8 @@ def _prediction_output_filename(grid_name: str, table_name: str) -> str: return f"{grid_name}_bus_embeddings.parquet" if table_name == "gen_embeddings": return f"{grid_name}_gen_embeddings.parquet" + if table_name == "branch": + return f"{grid_name}_branch_predictions.parquet" return f"{grid_name}_{table_name}_predictions.parquet" diff --git a/gridfm_graphkit/tasks/pf_task.py b/gridfm_graphkit/tasks/pf_task.py index e5b679af..13af5d59 100644 --- a/gridfm_graphkit/tasks/pf_task.py +++ b/gridfm_graphkit/tasks/pf_task.py @@ -26,6 +26,7 @@ plot_correlation_by_node_type, plot_residuals_histograms, residual_stats_by_type, + compute_branch_predictions, ) import torch import torch.distributed as dist @@ -495,7 +496,7 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0): mask_PV = batch.mask_dict["PV"] mask_REF = batch.mask_dict["REF"] - prediction_table = { + bus_predictions = { "scenario": scenario_ids.cpu().numpy(), "bus": local_bus_idx.cpu().numpy(), "Pd": bus_x[:, PD_H].cpu().numpy(), @@ -519,10 +520,23 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0): "reactive res. (MVar)": residual_Q.detach().cpu().numpy(), "PBE": residual_mva.detach().cpu().numpy(), } + + branch_predictions = compute_branch_predictions( + eval_bus, + target, + bus_edge_index, + bus_edge_attr, + scenario_ids, + local_bus_idx, + ) if embeddings is None or "bus" not in embeddings: - return prediction_table + return { + "bus": bus_predictions, + "branch": branch_predictions, + } return { - "bus": prediction_table, + "bus": bus_predictions, + "branch": branch_predictions, "bus_embeddings": embedding_table_from_tensor( embeddings["bus"], id_columns={ diff --git a/gridfm_graphkit/tasks/utils.py b/gridfm_graphkit/tasks/utils.py index 673ab365..69df3df7 100644 --- a/gridfm_graphkit/tasks/utils.py +++ b/gridfm_graphkit/tasks/utils.py @@ -1,9 +1,21 @@ import torch +import torch.nn.functional as F from torch_scatter import scatter_mean, scatter_max import matplotlib.pyplot as plt import seaborn as sns import numpy as np import os +from gridfm_graphkit.models.utils import ComputeBranchFlow +from gridfm_graphkit.datasets.globals import ( + VA_OUT, + ANG_MIN, + ANG_MAX, + RATE_A, + YFF_TT_R, + YFF_TT_I, + YFT_TF_R, + YFT_TF_I, +) def local_index_per_graph(batch_index: torch.Tensor) -> torch.Tensor: @@ -218,3 +230,85 @@ def plot_correlation_by_node_type( filename = f"{prefix}_correlation_{node_type}.png" plt.savefig(os.path.join(plot_dir, filename), dpi=300) plt.close(fig) + + +def compute_branch_predictions( + eval_bus, + target, + bus_edge_index, + bus_edge_attr, + scenario_ids, + local_bus_idx, +): + """Compute branch-level predictions and ground-truth constraint violations. + + Args: + eval_bus: Clamped model predictions [num_bus, 4]. Branch flows + and angle violations are computed from this. + target: Ground truth bus tensor [num_bus, 4]. Target branch + flows and angle violations are computed from this. + bus_edge_index: Edge index [2, num_edges] (batch-global bus indices). + bus_edge_attr: Edge features [num_edges, num_edge_features]. + scenario_ids: Scenario ID per bus [num_bus] (batch-global). + local_bus_idx: Per-graph local bus index [num_bus]. + + Returns: + dict of numpy arrays, one entry per directed edge. + """ + branch_flow_layer = ComputeBranchFlow() + + from_bus_idx = bus_edge_index[0] + to_bus_idx = bus_edge_index[1] + + # Branch limits — ANG_MIN/ANG_MAX restored to degrees by inverse_transform; + # convert to radians to match VA_OUT which stays in radians. + angle_min = bus_edge_attr[:, ANG_MIN] * torch.pi / 180.0 + angle_max = bus_edge_attr[:, ANG_MAX] * torch.pi / 180.0 + branch_thermal_limits = bus_edge_attr[:, RATE_A] + + def _branch_flows(bus_state): + Pft, Qft = branch_flow_layer(bus_state, bus_edge_index, bus_edge_attr) + Sft = torch.sqrt(Pft**2 + Qft**2) + thermal_excess = F.relu(Sft - branch_thermal_limits) + return Pft, Qft, thermal_excess + + def _angle_violations(bus_state): + angles = bus_state[:, VA_OUT] + diff = angles[from_bus_idx] - angles[to_bus_idx] + diff = (diff + torch.pi) % (2 * torch.pi) - torch.pi # wrap to [-pi, pi] + return diff, F.relu(angle_min - diff), F.relu(diff - angle_max) + + # Predicted + Pft, Qft, thermal_excess = _branch_flows(eval_bus) + angle_diff, angle_excess_low, angle_excess_high = _angle_violations(eval_bus) + + # Ground truth + Pft_target, Qft_target, thermal_excess_target = _branch_flows(target) + angle_diff_target, angle_excess_low_target, angle_excess_high_target = _angle_violations(target) + + def _np(t): + return t.detach().cpu().numpy() + + return { + "scenario": scenario_ids[from_bus_idx].cpu().numpy(), + "from_bus": local_bus_idx[from_bus_idx].cpu().numpy(), + "to_bus": local_bus_idx[to_bus_idx].cpu().numpy(), + "Pft": _np(Pft), + "Qft": _np(Qft), + "Pft_target": _np(Pft_target), + "Qft_target": _np(Qft_target), + "angle_diff": _np(angle_diff), + "angle_excess_low": _np(angle_excess_low), + "angle_excess_high": _np(angle_excess_high), + "angle_diff_target": _np(angle_diff_target), + "angle_excess_low_target": _np(angle_excess_low_target), + "angle_excess_high_target": _np(angle_excess_high_target), + "thermal_excess": _np(thermal_excess), + "thermal_excess_target": _np(thermal_excess_target), + # Fields needed for current-based loading computation + "rate_a": _np(branch_thermal_limits), + "Yff_r": _np(bus_edge_attr[:, YFF_TT_R]), + "Yff_i": _np(bus_edge_attr[:, YFF_TT_I]), + "Yft_r": _np(bus_edge_attr[:, YFT_TF_R]), + "Yft_i": _np(bus_edge_attr[:, YFT_TF_I]), + } From e12eda0d721ffc30b1c58616d949e67c6b81bfd6 Mon Sep 17 00:00:00 2001 From: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> Date: Fri, 18 Sep 2026 21:16:17 +0100 Subject: [PATCH 2/6] branch flow tests Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> --- tests/test_compute_branch_predictions.py | 238 +++++++++++++++++++++++ 1 file changed, 238 insertions(+) create mode 100644 tests/test_compute_branch_predictions.py diff --git a/tests/test_compute_branch_predictions.py b/tests/test_compute_branch_predictions.py new file mode 100644 index 00000000..41951649 --- /dev/null +++ b/tests/test_compute_branch_predictions.py @@ -0,0 +1,238 @@ +"""Tests for compute_branch_predictions in gridfm_graphkit.tasks.utils.""" + +import numpy as np +import pytest +import torch +import yaml +from torch_geometric.data import HeteroData + +from gridfm_graphkit.datasets.globals import ( + ANG_MAX, + ANG_MIN, + P_E, + Q_E, + RATE_A, + VM_H, + VA_H, + YFF_TT_I, + YFF_TT_R, + YFT_TF_I, + YFT_TF_R, +) +from gridfm_graphkit.datasets.normalizers import HeteroDataMVANormalizer +from gridfm_graphkit.io.param_handler import NestedNamespace +from gridfm_graphkit.tasks.utils import compute_branch_predictions + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +# edge_attr has 11 columns (indices 0–10); build a minimal valid tensor. +_NUM_EDGE_FEATURES = 11 + + +def _make_edge_attr( + num_edges: int, + *, + yff_r: float = 1.0, + yff_i: float = 0.0, + yft_r: float = -1.0, + yft_i: float = 0.0, + ang_min_deg: float = -30.0, + ang_max_deg: float = 30.0, + rate_a: float = 100.0, +) -> torch.Tensor: + """Return a [num_edges, 11] edge_attr tensor with controlled values.""" + attr = torch.zeros(num_edges, _NUM_EDGE_FEATURES) + attr[:, YFF_TT_R] = yff_r + attr[:, YFF_TT_I] = yff_i + attr[:, YFT_TF_R] = yft_r + attr[:, YFT_TF_I] = yft_i + attr[:, ANG_MIN] = ang_min_deg + attr[:, ANG_MAX] = ang_max_deg + attr[:, RATE_A] = rate_a + return attr + + +def _make_inputs( + num_bus: int = 3, + num_edges: int = 4, + *, + vm: float = 1.0, + va: float = 0.0, + rate_a: float = 100.0, + ang_min_deg: float = -30.0, + ang_max_deg: float = 30.0, +): + """Return (eval_bus, target, bus_edge_index, bus_edge_attr, scenario_ids, local_bus_idx).""" + # bus state: [VM, VA, PG, QG] + bus = torch.zeros(num_bus, 4) + bus[:, 0] = vm # VM_OUT = 0 + bus[:, 1] = va # VA_OUT = 1 + + # simple directed edges: 0->1, 1->2, 2->0, 1->0 + src = torch.tensor([0, 1, 2, 1])[:num_edges] + dst = torch.tensor([1, 2, 0, 0])[:num_edges] + edge_index = torch.stack([src, dst]) + + edge_attr = _make_edge_attr( + num_edges, + rate_a=rate_a, + ang_min_deg=ang_min_deg, + ang_max_deg=ang_max_deg, + ) + + scenario_ids = torch.zeros(num_bus, dtype=torch.long) + local_bus_idx = torch.arange(num_bus) + + return bus.clone(), bus.clone(), edge_index, edge_attr, scenario_ids, local_bus_idx + + +# --------------------------------------------------------------------------- +# Test 1 — output keys are exactly as expected +# --------------------------------------------------------------------------- + +EXPECTED_KEYS = { + "scenario", "from_bus", "to_bus", + "Pft", "Qft", "Pft_target", "Qft_target", + "angle_diff", "angle_excess_low", "angle_excess_high", + "angle_diff_target", "angle_excess_low_target", "angle_excess_high_target", + "thermal_excess", "thermal_excess_target", + "rate_a", "Yff_r", "Yff_i", "Yft_r", "Yft_i", +} + + +def test_output_keys(): + result = compute_branch_predictions(*_make_inputs()) + assert set(result.keys()) == EXPECTED_KEYS + + +# --------------------------------------------------------------------------- +# Test 2 — all values are numpy arrays +# --------------------------------------------------------------------------- + +def test_all_values_are_numpy(): + result = compute_branch_predictions(*_make_inputs()) + for key, val in result.items(): + assert isinstance(val, np.ndarray), f"{key!r} is {type(val)}, expected np.ndarray" + + +# --------------------------------------------------------------------------- +# Test 3 — all arrays have length == num_edges +# --------------------------------------------------------------------------- + +def test_array_lengths(): + num_edges = 4 + result = compute_branch_predictions(*_make_inputs(num_edges=num_edges)) + for key, val in result.items(): + assert len(val) == num_edges, f"{key!r} has length {len(val)}, expected {num_edges}" + + +# --------------------------------------------------------------------------- +# Test 4 — thermal excess is non-negative everywhere +# --------------------------------------------------------------------------- + +def test_thermal_excess_non_negative(): + result = compute_branch_predictions(*_make_inputs()) + assert (result["thermal_excess"] >= 0).all() + assert (result["thermal_excess_target"] >= 0).all() + + +# --------------------------------------------------------------------------- +# Test 5 — angle excess is non-negative everywhere +# --------------------------------------------------------------------------- + +def test_angle_excess_non_negative(): + result = compute_branch_predictions(*_make_inputs()) + assert (result["angle_excess_low"] >= 0).all() + assert (result["angle_excess_high"] >= 0).all() + assert (result["angle_excess_low_target"] >= 0).all() + assert (result["angle_excess_high_target"] >= 0).all() + + +# --------------------------------------------------------------------------- +# Test 6 — angle_diff is wrapped to [-pi, pi] +# --------------------------------------------------------------------------- + +def test_angle_diff_wrapped(): + result = compute_branch_predictions(*_make_inputs()) + assert (result["angle_diff"] >= -np.pi).all() + assert (result["angle_diff"] <= np.pi).all() + assert (result["angle_diff_target"] >= -np.pi).all() + assert (result["angle_diff_target"] <= np.pi).all() + + +# --------------------------------------------------------------------------- +# Test 7 — when eval_bus == target, predicted fields equal ground-truth fields +# --------------------------------------------------------------------------- + +def test_perfect_prediction_equals_target(): + args = _make_inputs() + result = compute_branch_predictions(*args) + np.testing.assert_array_equal(result["Pft"], result["Pft_target"]) + np.testing.assert_array_equal(result["Qft"], result["Qft_target"]) + np.testing.assert_array_equal(result["thermal_excess"], result["thermal_excess_target"]) + np.testing.assert_array_equal(result["angle_diff"], result["angle_diff_target"]) + np.testing.assert_array_equal(result["angle_excess_low"], result["angle_excess_low_target"]) + np.testing.assert_array_equal(result["angle_excess_high"], result["angle_excess_high_target"]) + + +# --------------------------------------------------------------------------- +# Test 8 — thermal excess is zero when apparent flow is well below rate_a +# --------------------------------------------------------------------------- + +def test_no_thermal_excess_when_within_limits(): + # VM=0.001 -> very small flows -> well below rate_a=100 + result = compute_branch_predictions(*_make_inputs(vm=0.001, rate_a=100.0)) + np.testing.assert_array_equal(result["thermal_excess"], 0.0) + np.testing.assert_array_equal(result["thermal_excess_target"], 0.0) + + +# --------------------------------------------------------------------------- +# Test 9 — branch flows match stored P_E/Q_E on real case14 data +# --------------------------------------------------------------------------- + +def test_branch_flows_match_stored_values(): + data_dict = torch.load( + "tests/data/case14_ieee/processed/data_index_0.pt", + weights_only=True, + ) + data = HeteroData.from_dict(data_dict) + + node_stats = torch.load( + "tests/data/case14_ieee/processed/data_stats_HeteroDataMVANormalizer.pt", + weights_only=True, + ) + with open("tests/config/datamodule_test_base_config.yaml", "r") as f: + args = NestedNamespace(**yaml.safe_load(f)) + + normalizer = HeteroDataMVANormalizer(args) + normalizer.fit_from_dict(node_stats) + normalizer.transform(data) + + bus_edge_index = data[("bus", "connects", "bus")].edge_index + bus_edge_attr = data[("bus", "connects", "bus")].edge_attr + num_bus = data["bus"].x.size(0) + + # Use ground-truth VM/VA as both eval_bus and target + bus_state = torch.zeros(num_bus, 4) + bus_state[:, 0] = data["bus"].x[:, VM_H] # VM_OUT = 0 + bus_state[:, 1] = data["bus"].x[:, VA_H] # VA_OUT = 1 + + scenario_ids = torch.zeros(num_bus, dtype=torch.long) + local_bus_idx = torch.arange(num_bus) + + result = compute_branch_predictions( + bus_state, + bus_state, + bus_edge_index, + bus_edge_attr, + scenario_ids, + local_bus_idx, + ) + + stored_P = bus_edge_attr[:, P_E].numpy() + stored_Q = bus_edge_attr[:, Q_E].numpy() + + np.testing.assert_allclose(result["Pft"], stored_P, atol=1e-4) + np.testing.assert_allclose(result["Qft"], stored_Q, atol=1e-4) From 4f8bdf332fbbb4761f48a021210f5b921a398513 Mon Sep 17 00:00:00 2001 From: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> Date: Fri, 18 Sep 2026 21:22:23 +0100 Subject: [PATCH 3/6] style Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> --- gridfm_graphkit/tasks/pf_task.py | 2 +- gridfm_graphkit/tasks/utils.py | 46 +++++++++--------- tests/test_compute_branch_predictions.py | 62 ++++++++++++++++++------ 3 files changed, 73 insertions(+), 37 deletions(-) diff --git a/gridfm_graphkit/tasks/pf_task.py b/gridfm_graphkit/tasks/pf_task.py index 13af5d59..a942c79a 100644 --- a/gridfm_graphkit/tasks/pf_task.py +++ b/gridfm_graphkit/tasks/pf_task.py @@ -528,7 +528,7 @@ def predict_step(self, batch, batch_idx, dataloader_idx=0): bus_edge_attr, scenario_ids, local_bus_idx, - ) + ) if embeddings is None or "bus" not in embeddings: return { "bus": bus_predictions, diff --git a/gridfm_graphkit/tasks/utils.py b/gridfm_graphkit/tasks/utils.py index 69df3df7..6db20040 100644 --- a/gridfm_graphkit/tasks/utils.py +++ b/gridfm_graphkit/tasks/utils.py @@ -258,7 +258,7 @@ def compute_branch_predictions( branch_flow_layer = ComputeBranchFlow() from_bus_idx = bus_edge_index[0] - to_bus_idx = bus_edge_index[1] + to_bus_idx = bus_edge_index[1] # Branch limits — ANG_MIN/ANG_MAX restored to degrees by inverse_transform; # convert to radians to match VA_OUT which stays in radians. @@ -284,31 +284,33 @@ def _angle_violations(bus_state): # Ground truth Pft_target, Qft_target, thermal_excess_target = _branch_flows(target) - angle_diff_target, angle_excess_low_target, angle_excess_high_target = _angle_violations(target) + angle_diff_target, angle_excess_low_target, angle_excess_high_target = ( + _angle_violations(target) + ) def _np(t): return t.detach().cpu().numpy() return { - "scenario": scenario_ids[from_bus_idx].cpu().numpy(), - "from_bus": local_bus_idx[from_bus_idx].cpu().numpy(), - "to_bus": local_bus_idx[to_bus_idx].cpu().numpy(), - "Pft": _np(Pft), - "Qft": _np(Qft), - "Pft_target": _np(Pft_target), - "Qft_target": _np(Qft_target), - "angle_diff": _np(angle_diff), - "angle_excess_low": _np(angle_excess_low), - "angle_excess_high": _np(angle_excess_high), - "angle_diff_target": _np(angle_diff_target), - "angle_excess_low_target": _np(angle_excess_low_target), - "angle_excess_high_target": _np(angle_excess_high_target), - "thermal_excess": _np(thermal_excess), - "thermal_excess_target": _np(thermal_excess_target), + "scenario": scenario_ids[from_bus_idx].cpu().numpy(), + "from_bus": local_bus_idx[from_bus_idx].cpu().numpy(), + "to_bus": local_bus_idx[to_bus_idx].cpu().numpy(), + "Pft": _np(Pft), + "Qft": _np(Qft), + "Pft_target": _np(Pft_target), + "Qft_target": _np(Qft_target), + "angle_diff": _np(angle_diff), + "angle_excess_low": _np(angle_excess_low), + "angle_excess_high": _np(angle_excess_high), + "angle_diff_target": _np(angle_diff_target), + "angle_excess_low_target": _np(angle_excess_low_target), + "angle_excess_high_target": _np(angle_excess_high_target), + "thermal_excess": _np(thermal_excess), + "thermal_excess_target": _np(thermal_excess_target), # Fields needed for current-based loading computation - "rate_a": _np(branch_thermal_limits), - "Yff_r": _np(bus_edge_attr[:, YFF_TT_R]), - "Yff_i": _np(bus_edge_attr[:, YFF_TT_I]), - "Yft_r": _np(bus_edge_attr[:, YFT_TF_R]), - "Yft_i": _np(bus_edge_attr[:, YFT_TF_I]), + "rate_a": _np(branch_thermal_limits), + "Yff_r": _np(bus_edge_attr[:, YFF_TT_R]), + "Yff_i": _np(bus_edge_attr[:, YFF_TT_I]), + "Yft_r": _np(bus_edge_attr[:, YFT_TF_R]), + "Yft_i": _np(bus_edge_attr[:, YFT_TF_I]), } diff --git a/tests/test_compute_branch_predictions.py b/tests/test_compute_branch_predictions.py index 41951649..3e89482b 100644 --- a/tests/test_compute_branch_predictions.py +++ b/tests/test_compute_branch_predictions.py @@ -1,7 +1,6 @@ """Tests for compute_branch_predictions in gridfm_graphkit.tasks.utils.""" import numpy as np -import pytest import torch import yaml from torch_geometric.data import HeteroData @@ -67,8 +66,8 @@ def _make_inputs( """Return (eval_bus, target, bus_edge_index, bus_edge_attr, scenario_ids, local_bus_idx).""" # bus state: [VM, VA, PG, QG] bus = torch.zeros(num_bus, 4) - bus[:, 0] = vm # VM_OUT = 0 - bus[:, 1] = va # VA_OUT = 1 + bus[:, 0] = vm # VM_OUT = 0 + bus[:, 1] = va # VA_OUT = 1 # simple directed edges: 0->1, 1->2, 2->0, 1->0 src = torch.tensor([0, 1, 2, 1])[:num_edges] @@ -93,12 +92,26 @@ def _make_inputs( # --------------------------------------------------------------------------- EXPECTED_KEYS = { - "scenario", "from_bus", "to_bus", - "Pft", "Qft", "Pft_target", "Qft_target", - "angle_diff", "angle_excess_low", "angle_excess_high", - "angle_diff_target", "angle_excess_low_target", "angle_excess_high_target", - "thermal_excess", "thermal_excess_target", - "rate_a", "Yff_r", "Yff_i", "Yft_r", "Yft_i", + "scenario", + "from_bus", + "to_bus", + "Pft", + "Qft", + "Pft_target", + "Qft_target", + "angle_diff", + "angle_excess_low", + "angle_excess_high", + "angle_diff_target", + "angle_excess_low_target", + "angle_excess_high_target", + "thermal_excess", + "thermal_excess_target", + "rate_a", + "Yff_r", + "Yff_i", + "Yft_r", + "Yft_i", } @@ -111,27 +124,34 @@ def test_output_keys(): # Test 2 — all values are numpy arrays # --------------------------------------------------------------------------- + def test_all_values_are_numpy(): result = compute_branch_predictions(*_make_inputs()) for key, val in result.items(): - assert isinstance(val, np.ndarray), f"{key!r} is {type(val)}, expected np.ndarray" + assert isinstance(val, np.ndarray), ( + f"{key!r} is {type(val)}, expected np.ndarray" + ) # --------------------------------------------------------------------------- # Test 3 — all arrays have length == num_edges # --------------------------------------------------------------------------- + def test_array_lengths(): num_edges = 4 result = compute_branch_predictions(*_make_inputs(num_edges=num_edges)) for key, val in result.items(): - assert len(val) == num_edges, f"{key!r} has length {len(val)}, expected {num_edges}" + assert len(val) == num_edges, ( + f"{key!r} has length {len(val)}, expected {num_edges}" + ) # --------------------------------------------------------------------------- # Test 4 — thermal excess is non-negative everywhere # --------------------------------------------------------------------------- + def test_thermal_excess_non_negative(): result = compute_branch_predictions(*_make_inputs()) assert (result["thermal_excess"] >= 0).all() @@ -142,6 +162,7 @@ def test_thermal_excess_non_negative(): # Test 5 — angle excess is non-negative everywhere # --------------------------------------------------------------------------- + def test_angle_excess_non_negative(): result = compute_branch_predictions(*_make_inputs()) assert (result["angle_excess_low"] >= 0).all() @@ -154,6 +175,7 @@ def test_angle_excess_non_negative(): # Test 6 — angle_diff is wrapped to [-pi, pi] # --------------------------------------------------------------------------- + def test_angle_diff_wrapped(): result = compute_branch_predictions(*_make_inputs()) assert (result["angle_diff"] >= -np.pi).all() @@ -166,21 +188,32 @@ def test_angle_diff_wrapped(): # Test 7 — when eval_bus == target, predicted fields equal ground-truth fields # --------------------------------------------------------------------------- + def test_perfect_prediction_equals_target(): args = _make_inputs() result = compute_branch_predictions(*args) np.testing.assert_array_equal(result["Pft"], result["Pft_target"]) np.testing.assert_array_equal(result["Qft"], result["Qft_target"]) - np.testing.assert_array_equal(result["thermal_excess"], result["thermal_excess_target"]) + np.testing.assert_array_equal( + result["thermal_excess"], + result["thermal_excess_target"], + ) np.testing.assert_array_equal(result["angle_diff"], result["angle_diff_target"]) - np.testing.assert_array_equal(result["angle_excess_low"], result["angle_excess_low_target"]) - np.testing.assert_array_equal(result["angle_excess_high"], result["angle_excess_high_target"]) + np.testing.assert_array_equal( + result["angle_excess_low"], + result["angle_excess_low_target"], + ) + np.testing.assert_array_equal( + result["angle_excess_high"], + result["angle_excess_high_target"], + ) # --------------------------------------------------------------------------- # Test 8 — thermal excess is zero when apparent flow is well below rate_a # --------------------------------------------------------------------------- + def test_no_thermal_excess_when_within_limits(): # VM=0.001 -> very small flows -> well below rate_a=100 result = compute_branch_predictions(*_make_inputs(vm=0.001, rate_a=100.0)) @@ -192,6 +225,7 @@ def test_no_thermal_excess_when_within_limits(): # Test 9 — branch flows match stored P_E/Q_E on real case14 data # --------------------------------------------------------------------------- + def test_branch_flows_match_stored_values(): data_dict = torch.load( "tests/data/case14_ieee/processed/data_index_0.pt", From 78f452afbc1080dc468ec10e1d2df0b6afa5144b Mon Sep 17 00:00:00 2001 From: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:17:00 +0100 Subject: [PATCH 4/6] docs Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> --- README.md | 4 ++-- docs/quick_start/quick_start.md | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index e506f290..954f118e 100644 --- a/README.md +++ b/README.md @@ -242,7 +242,7 @@ gridfm_graphkit evaluate --config path/to/eval.yaml --model_path path/to/model.p | `--dataset_wrapper_cache_dir` | `str` | Disk cache directory for dataset wrapper; cache is loaded from here when present and saved after first population. | `None` | | `--profiler` | `str` | Enable Lightning profiler (`simple`, `advanced`, `pytorch`). | `None` | | `--compute_dc_ac_metrics` | `flag` | Compute ground-truth AC/DC power balance metrics on the test split. | `False` | -| `--save_output` | `flag` | Save predictions as `_predictions.parquet` under MLflow artifacts (`.../artifacts/test`). | `False` | +| `--save_output` | `flag` | Save predictions under MLflow artifacts (`.../artifacts/test`). For the PowerFlow task this writes `_predictions.parquet` (bus-level) and `_branch_predictions.parquet` (branch-level flows, thermal loading, and angle violations). | `False` | | `--mp_context` | `str` | DataLoader multiprocessing start method (`spawn`, `fork`, `forkserver`). Defaults to PyTorch's automatic choice. On Linux, `spawn` is recommended for safety (CUDA + fork is unsafe); other choices emit a warning. | `None` | ### Example with saved normalizer stats @@ -282,7 +282,7 @@ gridfm_graphkit predict --config path/to/config.yaml --model_path path/to/model. | `--plugins` | `list[str]` | Python packages to import for plugin registration, e.g. `gridfm_graphkit_ee`. | `[]` | | `--num_workers` | `int` | Override `data.workers` from YAML. Use `0` to debug worker crashes. | `None` | | `--dataset_wrapper_cache_dir` | `str` | Disk cache directory for dataset wrapper; cache is loaded from here when present and saved after first population. | `None` | -| `--output_path` | `str` | Directory where predictions are saved as `_predictions.parquet`. | `data` | +| `--output_path` | `str` | Directory where predictions are saved. For the PowerFlow task this writes `_predictions.parquet` (bus-level) and `_branch_predictions.parquet` (branch-level flows, thermal loading, and angle violations). | `data` | | `--get_embeddings` | `flag` | Export final hidden embeddings to `_bus_embeddings.parquet` (and `_gen_embeddings.parquet` for OPF models that expose gen embeddings) in `--output_path`. | `False` | | `--compile [MODE]` | `str` | Enable `torch.compile` mode. Valid values: `default`, `reduce-overhead`, `max-autotune`, `max-autotune-no-cudagraphs`. If flag is passed without a value, mode is `default`. | `None` | | `--bfloat16` | `flag` | Cast model to `torch.bfloat16` (`model.to(torch.bfloat16)`). | `False` | diff --git a/docs/quick_start/quick_start.md b/docs/quick_start/quick_start.md index 2b1b63fc..5e619525 100644 --- a/docs/quick_start/quick_start.md +++ b/docs/quick_start/quick_start.md @@ -108,7 +108,7 @@ gridfm_graphkit evaluate --config path/to/eval.yaml --model_path path/to/model.p | `--dataset_wrapper_cache_dir` | `str` | Disk cache directory for dataset wrapper; cache is loaded from here when present and saved after first population. | `None` | | `--profiler` | `str` | Enable Lightning profiler (`simple`, `advanced`, `pytorch`). | `None` | | `--compute_dc_ac_metrics` | `flag` | Compute ground-truth AC/DC power balance metrics on the test split. | `False` | -| `--save_output` | `flag` | Save predictions as `_predictions.parquet` under MLflow artifacts (`.../artifacts/test`). | `False` | +| `--save_output` | `flag` | Save predictions under MLflow artifacts (`.../artifacts/test`). For the PowerFlow task this writes `_predictions.parquet` (bus-level) and `_branch_predictions.parquet` (branch-level flows, thermal loading, and angle violations). | `False` | | `--mp_context` | `str` | DataLoader multiprocessing start method (`spawn`, `fork`, `forkserver`). Defaults to PyTorch's automatic choice. On Linux, `spawn` is recommended for safety (CUDA + fork is unsafe); other choices emit a warning. | `None` | ### Example with saved normalizer stats @@ -148,7 +148,7 @@ gridfm_graphkit predict --config path/to/config.yaml --model_path path/to/model. | `--plugins` | `list[str]` | Python packages to import for plugin registration, e.g. `gridfm_graphkit_ee`. | `[]` | | `--num_workers` | `int` | Override `data.workers` from YAML. Use `0` to debug worker crashes. | `None` | | `--dataset_wrapper_cache_dir` | `str` | Disk cache directory for dataset wrapper; cache is loaded from here when present and saved after first population. | `None` | -| `--output_path` | `str` | Directory where predictions are saved as `_predictions.parquet`. | `data` | +| `--output_path` | `str` | Directory where predictions are saved. For the PowerFlow task this writes `_predictions.parquet` (bus-level) and `_branch_predictions.parquet` (branch-level flows, thermal loading, and angle violations). | `data` | | `--get_embeddings` | `flag` | Export final hidden embeddings to `_bus_embeddings.parquet` (and `_gen_embeddings.parquet` for OPF models that expose gen embeddings) in `--output_path`. | `False` | | `--compile [MODE]` | `str` | Enable `torch.compile` mode. Valid values: `default`, `reduce-overhead`, `max-autotune`, `max-autotune-no-cudagraphs`. If flag is passed without a value, mode is `default`. | `None` | | `--bfloat16` | `flag` | Cast model to `torch.bfloat16` (`model.to(torch.bfloat16)`). | `False` | From 9036725baa623b718a21ec25f43562b0b3b1e7fd Mon Sep 17 00:00:00 2001 From: naomi-simumba <7224231+naomi-simumba@users.noreply.github.com> Date: Wed, 23 Sep 2026 09:58:24 +0100 Subject: [PATCH 5/6] Update gridfm_graphkit/tasks/utils.py Co-authored-by: PUECH Alban <72336171+albanpuech@users.noreply.github.com> Signed-off-by: naomi-simumba <7224231+naomi-simumba@users.noreply.github.com> --- gridfm_graphkit/tasks/utils.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/gridfm_graphkit/tasks/utils.py b/gridfm_graphkit/tasks/utils.py index 6db20040..8e168819 100644 --- a/gridfm_graphkit/tasks/utils.py +++ b/gridfm_graphkit/tasks/utils.py @@ -241,7 +241,16 @@ def compute_branch_predictions( local_bus_idx, ): """Compute branch-level predictions and ground-truth constraint violations. - + + Expects tensors after ``inverse_transform``: ``Va`` stays in radians, while + ``ANG_MIN`` / ``ANG_MAX`` are restored to degrees and converted here with + ``* pi / 180``. Do not call this on training-space (post-``transform``) graphs: + those limits are already in radians, so the extra conversion would be wrong. + + Bidirectional edges reuse the same ``angmin`` / ``angmax``. That is correct + while limits are symmetric (e.g. ±30°). An asymmetric pair would be wrong + on the reverse copy. + Args: eval_bus: Clamped model predictions [num_bus, 4]. Branch flows and angle violations are computed from this. @@ -251,7 +260,7 @@ def compute_branch_predictions( bus_edge_attr: Edge features [num_edges, num_edge_features]. scenario_ids: Scenario ID per bus [num_bus] (batch-global). local_bus_idx: Per-graph local bus index [num_bus]. - + Returns: dict of numpy arrays, one entry per directed edge. """ From 78a788543d360a9767671155dee6ea30dcb645dc Mon Sep 17 00:00:00 2001 From: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> Date: Wed, 23 Sep 2026 10:12:03 +0100 Subject: [PATCH 6/6] style Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com> --- gridfm_graphkit/tasks/utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/gridfm_graphkit/tasks/utils.py b/gridfm_graphkit/tasks/utils.py index 503f3203..57ac4119 100644 --- a/gridfm_graphkit/tasks/utils.py +++ b/gridfm_graphkit/tasks/utils.py @@ -292,16 +292,16 @@ def compute_branch_predictions( local_bus_idx, ): """Compute branch-level predictions and ground-truth constraint violations. - + Expects tensors after ``inverse_transform``: ``Va`` stays in radians, while ``ANG_MIN`` / ``ANG_MAX`` are restored to degrees and converted here with ``* pi / 180``. Do not call this on training-space (post-``transform``) graphs: those limits are already in radians, so the extra conversion would be wrong. - + Bidirectional edges reuse the same ``angmin`` / ``angmax``. That is correct while limits are symmetric (e.g. ±30°). An asymmetric pair would be wrong on the reverse copy. - + Args: eval_bus: Clamped model predictions [num_bus, 4]. Branch flows and angle violations are computed from this. @@ -311,7 +311,7 @@ def compute_branch_predictions( bus_edge_attr: Edge features [num_edges, num_edge_features]. scenario_ids: Scenario ID per bus [num_bus] (batch-global). local_bus_idx: Per-graph local bus index [num_bus]. - + Returns: dict of numpy arrays, one entry per directed edge. """