Add branch-level outputs to PowerFlow predict_step

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 `<grid_name>_predictions.parquet` under MLflow artifacts (`.../artifacts/test`). | `False` |
+| `--save_output` | `flag` | Save predictions under MLflow artifacts (`.../artifacts/test`). For the PowerFlow task this writes `<grid_name>_predictions.parquet` (bus-level) and `<grid_name>_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 `<grid_name>_predictions.parquet`. | `data` |
+| `--output_path` | `str` | Directory where predictions are saved. For the PowerFlow task this writes `<grid_name>_predictions.parquet` (bus-level) and `<grid_name>_branch_predictions.parquet` (branch-level flows, thermal loading, and angle violations). | `data` |
 | `--get_embeddings` | `flag` | Export final hidden embeddings to `<grid_name>_bus_embeddings.parquet` (and `<grid_name>_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` |

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 `<grid_name>_predictions.parquet` under MLflow artifacts (`.../artifacts/test`). | `False` |
+| `--save_output` | `flag` | Save predictions under MLflow artifacts (`.../artifacts/test`). For the PowerFlow task this writes `<grid_name>_predictions.parquet` (bus-level) and `<grid_name>_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 `<grid_name>_predictions.parquet`. | `data` |
+| `--output_path` | `str` | Directory where predictions are saved. For the PowerFlow task this writes `<grid_name>_predictions.parquet` (bus-level) and `<grid_name>_branch_predictions.parquet` (branch-level flows, thermal loading, and angle violations). | `data` |
 | `--get_embeddings` | `flag` | Export final hidden embeddings to `<grid_name>_bus_embeddings.parquet` (and `<grid_name>_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` |

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"
 
 

gridfm_graphkit/tasks/pf_task.py

@@ -27,6 +27,7 @@
     plot_correlation_by_node_type,
     plot_residuals_histograms,
     residual_stats_by_type,
+    compute_branch_predictions,
 )
 import torch
 import torch.distributed as dist
@@ -498,7 +499,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(),
@@ -522,10 +523,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={

gridfm_graphkit/tasks/utils.py

@@ -5,6 +5,17 @@
 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 compute_angle_violation(
@@ -270,3 +281,96 @@ 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.
+
+    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.
+        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]),
+    }