diff --git a/src/power_grid_model_ds/_core/model/grids/_filter.py b/src/power_grid_model_ds/_core/model/grids/_filter.py new file mode 100644 index 00000000..8093eab9 --- /dev/null +++ b/src/power_grid_model_ds/_core/model/grids/_filter.py @@ -0,0 +1,209 @@ +# SPDX-FileCopyrightText: Contributors to the Power Grid Model project +# +# SPDX-License-Identifier: MPL-2.0 + +from typing import TYPE_CHECKING + +from power_grid_model_ds._core.model.constants import EMPTY_ID +from power_grid_model_ds._core.model.grids._search import get_downstream_nodes as _get_downstream_nodes + +if TYPE_CHECKING: + from power_grid_model_ds._core.model.grids.base import Grid + + +def filter_grid( + grid: "Grid", + feeder_ids: list[int], + *, + include_adjacent_nodes: bool = False, +) -> "Grid": + """Create a new Grid containing only the components belonging to the given feeders. + + All nodes and branches belonging to the specified feeder_ids are included, as well + as the feeding substation node for each feeder. All branch types whose endpoints are + both in the visible node set are included. All appliances (loads, generators, sources, + shunts) whose connected node is visible are included. Three-winding transformers are + included only when all three of their nodes are visible. + + Requires set_feeder_ids() to have been called on the grid first. + + Args: + grid: The source Grid to filter. + feeder_ids: IDs of feeder branches (feeder_branch_id values set by + set_feeder_ids()). All nodes and branches belonging to these feeders are + included, as well as the feeding substation node for each feeder. + include_adjacent_nodes: When True, also include nodes one hop away from any + visible node (via any branch or three-winding transformer). + + Returns: + A new Grid containing the filtered subset, with graphs rebuilt. + + Example: + >>> grid.set_feeder_ids() + >>> subset = filter_grid(grid, feeder_ids=[201]) + >>> subset = filter_grid(grid, feeder_ids=[201, 204]) + >>> subset = filter_grid(grid, feeder_ids=[201], include_adjacent_nodes=True) + """ + visible_nodes = _resolve_visible_nodes(grid, feeder_ids) + if include_adjacent_nodes: + visible_nodes = _expand_adjacent_nodes(grid, visible_nodes) + return _build_grid_subset(grid, visible_nodes) + + +def filter_path( + grid: "Grid", + start_node_id: int, + end_node_id: int, + *, + include_adjacent_nodes: bool = False, +) -> "Grid": + """Create a new Grid containing only the components along the shortest path between two nodes. + + Uses the active graph (active branches only) to find the shortest path. All nodes along + the path are included, as well as all branches connecting consecutive path nodes. + + Args: + grid: The source Grid to filter. + start_node_id: The starting node ID. + end_node_id: The ending node ID. + include_adjacent_nodes: When True, also include nodes one hop away from any + visible node (via any branch or three-winding transformer). + + Returns: + A new Grid containing the filtered subset, with graphs rebuilt. + + Raises: + NoPathBetweenNodes: if no path exists between the two nodes in the active graph. + + Example: + >>> subset = filter_path(grid, start_node_id=102, end_node_id=106) + >>> subset = filter_path(grid, start_node_id=102, end_node_id=106, include_adjacent_nodes=True) + """ + path_nodes, _ = grid.graphs.active_graph.get_shortest_path(start_node_id, end_node_id) + visible_nodes = set(path_nodes) + if include_adjacent_nodes: + visible_nodes = _expand_adjacent_nodes(grid, visible_nodes) + return _build_grid_subset(grid, visible_nodes) + + +def filter_nodes( + grid: "Grid", + node_ids: list[int], + *, + include_adjacent_nodes: bool = False, +) -> "Grid": + """Create a new Grid containing only the specified nodes and their connecting components. + + All branches whose both endpoints are in the given node list are included. All + appliances connected to any of the given nodes are included. Three-winding transformers + are included only when all three of their nodes are in the list. + + Args: + grid: The source Grid to filter. + node_ids: List of node IDs to include. + include_adjacent_nodes: When True, also include nodes one hop away from any + visible node (via any branch or three-winding transformer). + + Returns: + A new Grid containing the filtered subset, with graphs rebuilt. + + Example: + >>> subset = filter_nodes(grid, node_ids=[102, 103]) + >>> subset = filter_nodes(grid, node_ids=[102, 103], include_adjacent_nodes=True) + """ + visible_nodes = set(node_ids) + if include_adjacent_nodes: + visible_nodes = _expand_adjacent_nodes(grid, visible_nodes) + return _build_grid_subset(grid, visible_nodes) + + +def filter_downstream( + grid: "Grid", + node_id: int, + *, + include_adjacent_nodes: bool = False, +) -> "Grid": + """Create a new Grid containing only the nodes downstream of the given node. + + Uses the active graph and the grid's substation nodes to determine the downstream + direction. The given node itself is always included. + + Args: + grid: The source Grid to filter. + node_id: The node ID from which to start the downstream traversal. + include_adjacent_nodes: When True, also include nodes one hop away from any + visible node (including upstream-adjacent nodes). + + Returns: + A new Grid containing the filtered subset, with graphs rebuilt. + + Raises: + NotImplementedError: if node_id is a substation node. + + Example: + >>> subset = filter_downstream(grid, node_id=102) + >>> subset = filter_downstream(grid, node_id=102, include_adjacent_nodes=True) + """ + downstream = _get_downstream_nodes(grid, node_id=node_id, inclusive=True) + visible_nodes = set(downstream) + if include_adjacent_nodes: + visible_nodes = _expand_adjacent_nodes(grid, visible_nodes) + return _build_grid_subset(grid, visible_nodes) + + +def _resolve_visible_nodes(grid: "Grid", feeder_ids: list[int]) -> set[int]: + feeder_nodes = grid.node.filter(feeder_branch_id=feeder_ids) + visible: set[int] = set(feeder_nodes.id.tolist()) + substation_ids = {nid for nid in feeder_nodes.feeder_node_id.tolist() if nid != EMPTY_ID} + return visible | substation_ids + + +def _expand_adjacent_nodes(grid: "Grid", visible_nodes: set[int]) -> set[int]: + seed = list(visible_nodes) + expanded = set(visible_nodes) + + branches = grid.branches + expanded |= set(branches.filter(from_node=seed).to_node.tolist()) + expanded |= set(branches.filter(to_node=seed).from_node.tolist()) + + twt = grid.three_winding_transformer + for node_field in ("node_1", "node_2", "node_3"): + matches = twt.filter(**{node_field: seed}) # type: ignore[arg-type] + if matches.size: + for other_field in ("node_1", "node_2", "node_3"): + if other_field != node_field: + expanded |= set(getattr(matches, other_field).tolist()) + + return expanded + + +def _build_grid_subset(grid: "Grid", visible_nodes: set[int]) -> "Grid": + node_list = list(visible_nodes) + result = type(grid).empty() + + result.append(grid.node.filter(id=node_list)) + + for branch_array in grid.branch_arrays: + subset = branch_array.filter(from_node=node_list, to_node=node_list, mode_="AND") + if subset.size: + result.append(subset) + + twt_subset = grid.three_winding_transformer.filter( + node_1=node_list, node_2=node_list, node_3=node_list, mode_="AND" + ) + if twt_subset.size: + result.append(twt_subset) + + for appliance_array in ( + grid.sym_load, + grid.sym_gen, + grid.source, + grid.asym_load, + grid.asym_gen, + grid.shunt, + ): + subset = appliance_array.filter(node=node_list) + if subset.size: + result.append(subset) + + return result diff --git a/src/power_grid_model_ds/utils.py b/src/power_grid_model_ds/utils.py new file mode 100644 index 00000000..5e3bb91a --- /dev/null +++ b/src/power_grid_model_ds/utils.py @@ -0,0 +1,11 @@ +# SPDX-FileCopyrightText: Contributors to the Power Grid Model project +# +# SPDX-License-Identifier: MPL-2.0 +from power_grid_model_ds._core.model.grids._filter import filter_downstream, filter_grid, filter_nodes, filter_path + +__all__ = [ + "filter_downstream", + "filter_grid", + "filter_nodes", + "filter_path", +] diff --git a/tests/unit/model/grids/test_filter.py b/tests/unit/model/grids/test_filter.py new file mode 100644 index 00000000..6cfd88e4 --- /dev/null +++ b/tests/unit/model/grids/test_filter.py @@ -0,0 +1,212 @@ +# SPDX-FileCopyrightText: Contributors to the Power Grid Model project +# +# SPDX-License-Identifier: MPL-2.0 + +import pytest + +from power_grid_model_ds import Grid +from power_grid_model_ds.utils import filter_downstream, filter_grid, filter_nodes, filter_path +from tests.fixtures.grid_classes import ExtendedGrid +from tests.fixtures.grids import build_basic_grid + +# pylint: disable=missing-function-docstring + + +@pytest.fixture +def grid_with_feeders(basic_grid): + basic_grid.set_feeder_ids() + return basic_grid + + +def test_filter_returns_grid_instance(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[201]) + assert isinstance(result, Grid) + + +def test_filter_feeder_includes_all_feeder_nodes(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[201]) + visible = set(result.node.id.tolist()) + assert {102, 103, 106}.issubset(visible) + + +def test_filter_feeder_includes_substation_node(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[201]) + assert 101 in result.node.id.tolist() + + +def test_filter_feeder_excludes_nodes_from_other_feeders(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[201]) + visible = set(result.node.id.tolist()) + assert 104 not in visible + assert 105 not in visible + + +def test_filter_feeder_includes_connecting_branches(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[201]) + assert 201 in result.line.id.tolist() + assert 202 in result.line.id.tolist() + assert 301 in result.transformer.id.tolist() + assert 204 not in result.line.id.tolist() + assert 601 not in result.link.id.tolist() + + +def test_filter_feeder_includes_appliances(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[201]) + assert 401 in result.sym_load.id.tolist() + assert 402 in result.sym_load.id.tolist() + assert 501 in result.source.id.tolist() + assert 403 not in result.sym_load.id.tolist() + assert 404 not in result.sym_load.id.tolist() + + +def test_filter_feeder_rebuilds_graphs(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[201]) + assert result.graphs is not None + assert result.graphs.complete_graph.external_to_internal(101) is not None + assert result.graphs.complete_graph.external_to_internal(102) is not None + + +def test_filter_multiple_feeders(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[201, 204]) + assert set(result.node.id.tolist()) == {101, 102, 103, 104, 105, 106} + + +def test_filter_adjacent_nodes_expands_one_hop(grid_with_feeders): + # Feeder 204 covers nodes 104, 105 + substation 101 + # With adjacent: also pull in nodes connected to 101, 104, 105 + result = filter_grid(grid_with_feeders, feeder_ids=[204], include_adjacent_nodes=True) + visible = set(result.node.id.tolist()) + # Node 102 is adjacent to 101 (via line 201) → included + assert 102 in visible + # Node 104 and 105 are two hops from substation via feeder 201, but directly adjacent to 601 + assert 104 in visible + assert 105 in visible + + +def test_filter_adjacent_nodes_includes_connecting_branches(grid_with_feeders): + result = filter_grid(grid_with_feeders, feeder_ids=[204], include_adjacent_nodes=True) + # Link 601 (104↔105) is within feeder 204 + assert 601 in result.link.id.tolist() + # Line 201 (101↔102) becomes visible because 101 is in feeder 204 and 102 is one hop away + assert 201 in result.line.id.tolist() + + +def test_filter_preserves_grid_subclass(): + extended = build_basic_grid(ExtendedGrid.empty()) + extended.set_feeder_ids() + result = filter_grid(extended, feeder_ids=[201]) + assert type(result) is ExtendedGrid + + +# --------------------------------------------------------------------------- +# filter_path +# --------------------------------------------------------------------------- + + +def test_filter_path_returns_grid_instance(basic_grid): + result = filter_path(basic_grid, start_node_id=102, end_node_id=106) + assert isinstance(result, Grid) + + +def test_filter_path_direct_connection(basic_grid): + # Transformer 301 connects 102 → 106 directly + result = filter_path(basic_grid, start_node_id=102, end_node_id=106) + assert set(result.node.id.tolist()) == {102, 106} + assert 301 in result.transformer.id.tolist() + + +def test_filter_path_multi_hop(basic_grid): + # 103-104 requires crossing the inactive gap: 103-202-102-201-101-204-105-601-104 + result = filter_path(basic_grid, start_node_id=103, end_node_id=104) + assert {101, 102, 103, 104, 105}.issubset(set(result.node.id.tolist())) + + +def test_filter_path_same_node(basic_grid): + result = filter_path(basic_grid, start_node_id=102, end_node_id=102) + assert set(result.node.id.tolist()) == {102} + + +def test_filter_path_adjacent_nodes_expands_one_hop(basic_grid): + # Path 102→103; adjacent should pull in 101 (via 201) and 106 (via transformer 301) + result = filter_path(basic_grid, start_node_id=102, end_node_id=103, include_adjacent_nodes=True) + visible = set(result.node.id.tolist()) + assert 101 in visible + assert 106 in visible + + +# --------------------------------------------------------------------------- +# filter_nodes +# --------------------------------------------------------------------------- + + +def test_filter_nodes_returns_grid_instance(basic_grid): + result = filter_nodes(basic_grid, node_ids=[102, 103]) + assert isinstance(result, Grid) + + +def test_filter_nodes_includes_only_specified_nodes(basic_grid): + result = filter_nodes(basic_grid, node_ids=[102, 103]) + assert set(result.node.id.tolist()) == {102, 103} + + +def test_filter_nodes_includes_connecting_branches(basic_grid): + result = filter_nodes(basic_grid, node_ids=[102, 103]) + assert 202 in result.line.id.tolist() + + +def test_filter_nodes_excludes_external_branches(basic_grid): + # Line 201 connects 101↔102; 101 is not in the list + result = filter_nodes(basic_grid, node_ids=[102, 103]) + assert 201 not in result.line.id.tolist() + assert 301 not in result.transformer.id.tolist() + + +def test_filter_nodes_includes_appliances(basic_grid): + result = filter_nodes(basic_grid, node_ids=[102, 103]) + assert 401 in result.sym_load.id.tolist() + assert 402 in result.sym_load.id.tolist() + assert 403 not in result.sym_load.id.tolist() + + +def test_filter_nodes_adjacent_nodes_expands_one_hop(basic_grid): + # Starting from {102}: adjacent via 201→101, via 202→103, via transformer 301→106 + result = filter_nodes(basic_grid, node_ids=[102], include_adjacent_nodes=True) + visible = set(result.node.id.tolist()) + assert 101 in visible + assert 103 in visible + assert 106 in visible + + +# --------------------------------------------------------------------------- +# filter_downstream +# --------------------------------------------------------------------------- + + +def test_filter_downstream_returns_grid_instance(basic_grid): + result = filter_downstream(basic_grid, node_id=102) + assert isinstance(result, Grid) + + +def test_filter_downstream_includes_node_itself(basic_grid): + result = filter_downstream(basic_grid, node_id=102) + assert 102 in result.node.id.tolist() + + +def test_filter_downstream_includes_downstream_nodes(basic_grid): + # Downstream of 102 (away from substation 101): 102, 103, 106 + result = filter_downstream(basic_grid, node_id=102) + assert {102, 103, 106}.issubset(set(result.node.id.tolist())) + + +def test_filter_downstream_excludes_upstream_nodes(basic_grid): + result = filter_downstream(basic_grid, node_id=102) + visible = set(result.node.id.tolist()) + assert 101 not in visible + assert 104 not in visible + assert 105 not in visible + + +def test_filter_downstream_adjacent_nodes_adds_upstream(basic_grid): + # Downstream of 102 = {102, 103, 106}; adjacent expansion adds 101 (one hop upstream) + result = filter_downstream(basic_grid, node_id=102, include_adjacent_nodes=True) + assert 101 in result.node.id.tolist()