Skip to content

Add branch-level outputs to PowerFlow predict_step - #122

Merged
romeokienzler merged 7 commits into
gridfm:mainfrom
naomi-simumba:new_branch_output
Sep 23, 2026
Merged

romeokienzler merged 7 commits into
gridfm:mainfrom
naomi-simumba:new_branch_output

Conversation

@naomi-simumba

@naomi-simumba naomi-simumba commented Sep 18, 2026 •

Copy link
Copy Markdown
Member

Summary

Closes [https://github.com//issues/96]. Extends the PowerFlow task's predict_step to return branch-level predictions alongside bus-level predictions, and wires them through the existing CLI save logic.

Changes

gridfm_graphkit/tasks/utils.py: Added compute_branch_predictions() — computes active/reactive branch flows (Pft, Qft), thermal excess, angle differences, and angle bound violations for both predicted and ground-truth bus states

gridfm_graphkit/tasks/pf_task.py: Imported compute_branch_predictions from tasks.utils. Called compute_branch_predictions in predict_step and added "branch" to the returned dict

gridfm_graphkit/cli.py: Added an explicit "branch" case to _prediction_output_filename so branch predictions are saved as {grid_name}_branch_predictions.parquet

tests/test_compute_branch_predictions.py: Unit and integration tests for compute_branch_predictions covering output structure, numerical properties, and correctness against the case14_ieee reference data

Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com>
Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com>
Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com>
@romeokienzler

Copy link
Copy Markdown
Collaborator

@naomi-simumba Thanks for this — a clean, well-scoped addition, and the dedicated test file (structure, numerical properties, and correctness against real case14 data) is exactly the kind of coverage we like to see. All CI checks are green, commits are signed off (DCO passing), and the code correctly reuses ComputeBranchFlow and the existing global column indices.

What's needed (minor):

  • Docs mention the new output file. README (--save_output / --output_path rows) and docs/quick_start/quick_start.md currently only describe <grid_name>_predictions.parquet. Since predict_step now emits a branch table saved as <grid_name>_branch_predictions.parquet, a short note there would help users discover it (CONTRIBUTING asks to update docs/README when relevant).

Everything else looks in good shape from a pre-review standpoint — no new params or dependencies to wire into the config YAMLs or pyproject.toml. The PR shows as BLOCKED only because it still needs a maintainer's review, not a branch issue. I'll defer the merge decision to a human maintainer.

— 🤖 _automated pre-review; a maintainer will follow up_

@albanpuech

Copy link
Copy Markdown
Collaborator

Looks good to me. README and docs/quick_start/quick_start.md still only mention <grid_name>_predictions.parquet. Please add <grid_name>_branch_predictions.parquet next to the --save_output / --output_path notes so users can find the new table.

@romeokienzler

Copy link
Copy Markdown
Collaborator

@albanpuech agreed — that's the one remaining item on this PR, and I can point @naomi-simumba at the exact spots:

  • README.md lines 245 (--save_output) and 285 (--output_path)
  • docs/quick_start/quick_start.md lines 111 (--save_output) and 151 (--output_path)

Each currently documents only <grid_name>_predictions.parquet; adding a note that predict_step also emits <grid_name>_branch_predictions.parquet (following the pattern already used for the _bus_embeddings.parquet / _gen_embeddings.parquet rows just below) covers it.

Nothing else outstanding from the pre-review: CI green, DCO signed off, no new params or deps to wire into the config YAMLs or pyproject.toml. Once the docs note lands this should be ready for a maintainer's final look.

— 🤖 _automated pre-review; a maintainer will follow up_

Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com>
@romeokienzler

Copy link
Copy Markdown
Collaborator

@naomi-simumba The new docs commit (78f452a) resolves the one remaining item — both the --save_output and --output_path rows in README.md and docs/quick_start/quick_start.md now document <grid_name>_branch_predictions.parquet alongside the bus-level table, and the added note on branch flows / thermal loading / angle violations is a nice touch.

That was the only outstanding point from the pre-review. CI is currently re-running on the new commit; assuming it comes back green (DCO already passes), this looks ready for a maintainer's final look. Deferring the merge decision to a human maintainer.

— 🤖 _automated pre-review; a maintainer will follow up_

@albanpuech albanpuech left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Optional: predict_step already runs ComputeBranchFlow for residuals. This helper runs it again on the same eval_bus. You could pass those Pft/Qft in instead of recomputing. Up to you.

Comment thread gridfm_graphkit/tasks/utils.py
naomi-simumba and others added 2 commits September 23, 2026 09:58
Co-authored-by: PUECH Alban <72336171+albanpuech@users.noreply.github.com>
Signed-off-by: naomi-simumba <7224231+naomi-simumba@users.noreply.github.com>
@romeokienzler

Copy link
Copy Markdown
Collaborator

@naomi-simumba The new docstring commit (9036725) is a nice clarification — spelling out that compute_branch_predictions expects post-inverse_transform tensors (radians vs. degrees for the angle limits) and that the symmetric bidirectional-limit assumption holds for ±30° pairs is exactly the kind of caveat future readers will thank you for.

One small heads-up: the three new blank lines inside the docstring were added with trailing whitespace ( ), and this repo runs the trailing-whitespace pre-commit hook. The pre-commit-run check is still pending, but it will likely flag those — running pre-commit run --all-files locally (or just stripping the whitespace on those blank lines) should clear it. Everything else I raised earlier is resolved; once pre-commit is green this looks ready for a maintainer's final look.

— 🤖 _automated pre-review; a maintainer will follow up_

Signed-off-by: Naomi Simumba <7224231+naomi-simumba@users.noreply.github.com>
@romeokienzler

Copy link
Copy Markdown
Collaborator

@naomi-simumba The new style commit (78a7885) strips the trailing whitespace from the docstring blank lines — that was the last thing I'd flagged, so pre-commit should now come back clean. Everything else raised earlier (docs for the new _branch_predictions.parquet table, the inverse_transform/radians docstring caveat) is already resolved.

CI is re-running on the new commit (pre-commit-run, pytests, security still pending; DCO already green). Assuming those pass, this looks ready for a maintainer's final look. Deferring the merge decision to a human maintainer.

— 🤖 _automated pre-review; a maintainer will follow up_

@romeokienzler romeokienzler left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

All CI checks green, DCO signed off, docs updated (both README and quick_start) to document _branch_predictions.parquet. The outstanding doc gap raised by @albanpuech and in prior pre-reviews is resolved. The optional refactor note (avoid recomputing branch flows) was marked non-blocking. Code is clean, tests are thorough. Approving.

@romeokienzler
romeokienzler merged commit ba57e19 into gridfm:main Sep 23, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants