power-grid-model-ds 1.0.13__py3-none-any.whl → 1.1.1__py3-none-any.whl

power_grid_model_ds/_core/model/graphs/models/base.py

@@ -3,6 +3,8 @@
 # SPDX-License-Identifier: MPL-2.0
 
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from typing import Generator
 
 import numpy as np
 from numpy._typing import NDArray
@@ -34,6 +36,14 @@ class BaseGraphModel(ABC):
     def nr_branches(self):
         """Returns the number of branches in the graph"""
 
+    @property
+    def all_branches(self) -> Generator[tuple[int, int], None, None]:
+        """Returns all branches in the graph."""
+        return (
+            (self.internal_to_external(source), self.internal_to_external(target))
+            for source, target in self._all_branches()
+        )
+
     @abstractmethod
     def external_to_internal(self, ext_node_id: int) -> int:
         """Convert external node id to internal node id (internal)
@@ -63,6 +73,14 @@ class BaseGraphModel(ABC):
 
         return self._has_node(node_id=internal_node_id)
 
+    def in_branches(self, node_id: int) -> Generator[tuple[int, int], None, None]:
+        """Return all branches that have the node as an endpoint."""
+        int_node_id = self.external_to_internal(node_id)
+        internal_edges = self._in_branches(int_node_id=int_node_id)
+        return (
+            (self.internal_to_external(source), self.internal_to_external(target)) for source, target in internal_edges
+        )
+
     def add_node(self, ext_node_id: int, raise_on_fail: bool = True) -> None:
         """Add a node to the graph."""
         if self.has_node(ext_node_id):
@@ -164,6 +182,28 @@ class BaseGraphModel(ABC):
             branches = _get_branch3_branches(branch3)
             self.delete_branch_array(branches, raise_on_fail=raise_on_fail)
 
+    @contextmanager
+    def tmp_remove_nodes(self, nodes: list[int]) -> Generator:
+        """Context manager that temporarily removes nodes and their branches from the graph.
+        Example:
+            >>> with graph.tmp_remove_nodes([1, 2, 3]):
+            >>>    assert not graph.has_node(1)
+            >>> assert graph.has_node(1)
+        In practice, this is useful when you want to e.g. calculate the shortest path between two nodes without
+        considering certain nodes.
+        """
+        edge_list = []
+        for node in nodes:
+            edge_list += list(self.in_branches(node))
+            self.delete_node(node)
+
+        yield
+
+        for node in nodes:
+            self.add_node(node)
+        for source, target in edge_list:
+            self.add_branch(source, target)
+
     def get_shortest_path(self, ext_start_node_id: int, ext_end_node_id: int) -> tuple[list[int], int]:
         """Calculate the shortest path between two nodes
 
@@ -235,8 +275,49 @@ class BaseGraphModel(ABC):
             nodes_to_ignore=self._externals_to_internals(nodes_to_ignore),
             inclusive=inclusive,
         )
+
         return self._internals_to_externals(nodes)
 
+    def find_first_connected(self, node_id: int, candidate_node_ids: list[int]) -> int:
+        """Find the first connected node to the node_id from the candidate_node_ids
+
+        Note:
+            If multiple candidate nodes are connected to the node, the first one found is returned.
+            There is no guarantee that the same candidate node will be returned each time.
+
+        Raises:
+            MissingNodeError: if no connected node is found
+            ValueError: if the node_id is in candidate_node_ids
+        """
+        internal_node_id = self.external_to_internal(node_id)
+        internal_candidates = self._externals_to_internals(candidate_node_ids)
+        if internal_node_id in internal_candidates:
+            raise ValueError("node_id cannot be in candidate_node_ids")
+        return self.internal_to_external(self._find_first_connected(internal_node_id, internal_candidates))
+
+    def get_downstream_nodes(self, node_id: int, start_node_ids: list[int], inclusive: bool = False) -> list[int]:
+        """Find all nodes downstream of the node_id with respect to the start_node_ids
+
+        Example:
+            given this graph: [1] - [2] - [3] - [4]
+            >>> graph.get_downstream_nodes(2, [1]) == [3, 4]
+            >>> graph.get_downstream_nodes(2, [1], inclusive=True) == [2, 3, 4]
+
+        args:
+            node_id: node id to start the search from
+            start_node_ids: list of node ids considered 'above' the node_id
+            inclusive: whether to include the given node id in the result
+        returns:
+            list of node ids sorted by distance, downstream of to the node id
+        """
+        connected_node = self.find_first_connected(node_id, start_node_ids)
+        path, _ = self.get_shortest_path(node_id, connected_node)
+        _, upstream_node, *_ = (
+            path  # path is at least 2 elements long or find_first_connected would have raised an error
+        )
+
+        return self.get_connected(node_id, [upstream_node], inclusive)
+
     def find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph.
         Returns:
@@ -270,9 +351,15 @@ class BaseGraphModel(ABC):
             return branch.is_active.item()
         return True
 
+    @abstractmethod
+    def _in_branches(self, int_node_id: int) -> Generator[tuple[int, int], None, None]: ...
+
     @abstractmethod
     def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bool = False) -> list[int]: ...
 
+    @abstractmethod
+    def _find_first_connected(self, node_id: int, candidate_node_ids: list[int]) -> int: ...
+
     @abstractmethod
     def _has_branch(self, from_node_id, to_node_id) -> bool: ...
 
@@ -307,6 +394,9 @@ class BaseGraphModel(ABC):
     @abstractmethod
     def _find_fundamental_cycles(self) -> list[list[int]]: ...
 
+    @abstractmethod
+    def _all_branches(self) -> Generator[tuple[int, int], None, None]: ...
+
 
 def _get_branch3_branches(branch3: Branch3Array) -> BranchArray:
     node_1 = branch3.node_1.item()

power_grid_model_ds/_core/model/graphs/models/rustworkx.py

@@ -3,10 +3,11 @@
 # SPDX-License-Identifier: MPL-2.0
 
 import logging
+from typing import Generator
 
 import rustworkx as rx
 from rustworkx import NoEdgeBetweenNodes
-from rustworkx.visit import BFSVisitor, PruneSearch
+from rustworkx.visit import BFSVisitor, PruneSearch, StopSearch
 
 from power_grid_model_ds._core.model.graphs.errors import MissingBranchError, MissingNodeError, NoPathBetweenNodes
 from power_grid_model_ds._core.model.graphs.models._rustworkx_search import find_fundamental_cycles_rustworkx
@@ -99,6 +100,16 @@ class RustworkxGraphModel(BaseGraphModel):
 
         return connected_nodes
 
+    def _in_branches(self, int_node_id: int) -> Generator[tuple[int, int], None, None]:
+        return ((source, target) for source, target, _ in self._graph.in_edges(int_node_id))
+
+    def _find_first_connected(self, node_id: int, candidate_node_ids: list[int]) -> int:
+        visitor = _NodeFinder(candidate_nodes=candidate_node_ids)
+        rx.bfs_search(self._graph, [node_id], visitor)
+        if visitor.found_node is None:
+            raise MissingNodeError(f"node {node_id} is not connected to any of the candidate nodes")
+        return visitor.found_node
+
     def _find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph using Rustworkx.
 
@@ -107,6 +118,9 @@ class RustworkxGraphModel(BaseGraphModel):
         """
         return find_fundamental_cycles_rustworkx(self._graph)
 
+    def _all_branches(self) -> Generator[tuple[int, int], None, None]:
+        return ((source, target) for source, target in self._graph.edge_list())
+
 
 class _NodeVisitor(BFSVisitor):
     def __init__(self, nodes_to_ignore: list[int]):
@@ -117,3 +131,16 @@ class _NodeVisitor(BFSVisitor):
         if v in self.nodes_to_ignore:
             raise PruneSearch
         self.nodes.append(v)
+
+
+class _NodeFinder(BFSVisitor):
+    """Visitor that stops the search when a candidate node is found"""
+
+    def __init__(self, candidate_nodes: list[int]):
+        self.candidate_nodes = candidate_nodes
+        self.found_node: int | None = None
+
+    def discover_vertex(self, v):
+        if v in self.candidate_nodes:
+            self.found_node = v
+            raise StopSearch

power_grid_model_ds/_core/model/grids/_text_sources.py

@@ -5,7 +5,6 @@
 """Create a grid from text a text file"""
 
 import logging
-from pathlib import Path
 from typing import TYPE_CHECKING
 
 from power_grid_model_ds._core.model.enums.nodes import NodeType
@@ -35,29 +34,28 @@ class TextSource:
     def __init__(self, grid_class: type["Grid"]):
         self.grid = grid_class.empty()
 
-    def load_grid_from_path(self, path: Path):
-        """Load assets from text file & sort them by id so that
-        they are ready to be appended to the grid"""
+    def load_from_txt(self, *args: str) -> "Grid":
+        """Load a grid from text"""
 
-        txt_nodes, txt_branches = self.read_txt(path)
+        text_lines = [line for arg in args for line in arg.strip().split("\n")]
+
+        txt_nodes, txt_branches = self.read_txt(text_lines)
         self.add_nodes(txt_nodes)
         self.add_branches(txt_branches)
         self.grid.set_feeder_ids()
         return self.grid
 
     @staticmethod
-    def read_txt(path: Path) -> tuple[set, dict]:
+    def read_txt(txt_lines: list[str]) -> tuple[set, dict]:
         """Extract assets from text"""
-        with open(path, "r", encoding="utf-8") as f:
-            txt_rows = f.readlines()
 
         txt_nodes = set()
         txt_branches = {}
-        for text_line in txt_rows:
+        for text_line in txt_lines:
             if not text_line.strip() or text_line.startswith("#"):
                 continue  # skip empty lines and comments
             try:
-                from_node_str, to_node_str, *comments = text_line.strip().split(" ")
+                from_node_str, to_node_str, *comments = text_line.strip().split()
             except ValueError as err:
                 raise ValueError(f"Text line '{text_line}' is invalid. Skipping...") from err
             comments = comments[0].split(",") if comments else []
@@ -116,4 +114,4 @@ class TextSource:
             new_branch.to_status = 0
         else:
             new_branch.to_status = 1
-        self.grid.append(new_branch)
+        self.grid.append(new_branch, check_max_id=False)

power_grid_model_ds/_core/model/grids/base.py

@@ -331,6 +331,7 @@ class Grid(FancyArrayContainer):
 
     def get_downstream_nodes(self, node_id: int, inclusive: bool = False):
         """Get the downstream nodes from a node.
+        Assuming each node has a single feeding substation and the grid is radial
 
         Example:
             given this graph: [1] - [2] - [3] - [4], with 1 being a substation node
@@ -349,15 +350,14 @@ class Grid(FancyArrayContainer):
         Returns:
             list[int]: The downstream nodes.
         """
-        substation_node_id = self.get_nearest_substation_node(node_id).id.item()
+        substation_nodes = self.node.filter(node_type=NodeType.SUBSTATION_NODE.value)
 
-        if node_id == substation_node_id:
+        if node_id in substation_nodes.id:
             raise NotImplementedError("get_downstream_nodes is not implemented for substation nodes!")
 
-        path_to_substation, _ = self.graphs.active_graph.get_shortest_path(node_id, substation_node_id)
-        upstream_node = path_to_substation[1]
-
-        return self.graphs.active_graph.get_connected(node_id, nodes_to_ignore=[upstream_node], inclusive=inclusive)
+        return self.graphs.active_graph.get_downstream_nodes(
+            node_id=node_id, start_node_ids=list(substation_nodes.id), inclusive=inclusive
+        )
 
     def cache(self, cache_dir: Path, cache_name: str, compress: bool = True):
         """Cache Grid to a folder
@@ -408,6 +408,21 @@ class Grid(FancyArrayContainer):
             raise TypeError(f"{pickle_path.name} is not a valid {cls.__name__} cache.")
         return grid
 
+    @classmethod
+    def from_txt(cls, *args: str):
+        """Build a grid from a list of strings
+
+        See the documentation for the expected format of the txt_lines
+
+        Args:
+            *args (str): The lines of the grid
+
+        Examples:
+            >>> Grid.from_txt("1 2", "2 3", "3 4 transformer", "4 5", "S1 6")
+            alternative: Grid.from_txt("1 2\n2 3\n3 4 transformer\n4 5\nS1 6")
+        """
+        return TextSource(grid_class=cls).load_from_txt(*args)
+
     @classmethod
     # pylint: disable=arguments-differ
     def from_txt_file(cls, txt_file_path: Path):
@@ -416,8 +431,9 @@ class Grid(FancyArrayContainer):
         Args:
             txt_file_path (Path): The path to the txt file
         """
-        text_source = TextSource(grid_class=cls)
-        return text_source.load_grid_from_path(txt_file_path)
+        with open(txt_file_path, "r", encoding="utf-8") as f:
+            txt_lines = f.readlines()
+        return TextSource(grid_class=cls).load_from_txt(*txt_lines)
 
     def set_feeder_ids(self):
         """Sets feeder and substation id properties in the grids arrays"""

{power_grid_model_ds-1.0.13.dist-info → power_grid_model_ds-1.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: power-grid-model-ds
-Version: 1.0.13
+Version: 1.1.1
 Summary: Power Grid Model extension which provides a grid data structure for simulation and analysis
 Author-email: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
 License: MPL-2.0

{power_grid_model_ds-1.0.13.dist-info → power_grid_model_ds-1.1.1.dist-info}/RECORD

@@ -47,18 +47,18 @@ power_grid_model_ds/_core/model/graphs/container.py,sha256=w3nwfR5sM-EqVPmZWSmhI
 power_grid_model_ds/_core/model/graphs/errors.py,sha256=-9fAsCqpKM0-lCM-tO2fVot9hqZBJa8WCxxqWEH_z6k,479
 power_grid_model_ds/_core/model/graphs/models/__init__.py,sha256=xwhMajQKIAfL1R3cq2sMNXMWtqfI5WGBi6-P4clyTes,262
 power_grid_model_ds/_core/model/graphs/models/_rustworkx_search.py,sha256=TqRR22TPx8XjRLukQxK8R8FjZvJ60CttDFctbISSbOM,2979
-power_grid_model_ds/_core/model/graphs/models/base.py,sha256=5IzmRGyacpoFNI09GEDGV383pBhtMIQfklTuRw1Utpc,12792
-power_grid_model_ds/_core/model/graphs/models/rustworkx.py,sha256=jD0Bn2tXYdnTKEUPc-eh-sie6o7ME5ukf62VwNSeUUc,4589
+power_grid_model_ds/_core/model/graphs/models/base.py,sha256=cHwwQFgPW78P5JDyfGE1fXEfwG7cv9-aG05UKPd_vGo,16710
+power_grid_model_ds/_core/model/graphs/models/rustworkx.py,sha256=a7PrDgge2KzMbeTQhQTE-KRPdKNj5TJvfYkzYNT8FtM,5730
 power_grid_model_ds/_core/model/grids/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-power_grid_model_ds/_core/model/grids/_text_sources.py,sha256=Ywt2c1YzBv88w778pr33MVJeCYQX51TjwCgL70M3MPw,4306
-power_grid_model_ds/_core/model/grids/base.py,sha256=FnUR0112xHvmKwY9u3_yTMARlO8btXGPj9dmNUZMJUM,16074
+power_grid_model_ds/_core/model/grids/_text_sources.py,sha256=3CWn0zjaUkFhrnSEIqAas52xeW0EX3EWXIXyo8XSQgw,4229
+power_grid_model_ds/_core/model/grids/base.py,sha256=knOYFl5snBFjccpbDPgWMrgCaFYN6JS74L9O15_snAI,16599
 power_grid_model_ds/_core/model/grids/helpers.py,sha256=VmIQ5a8oBMWW-VT5HOJk4iRRfEyO_OKc9LqKG0eVpSs,4174
 power_grid_model_ds/_core/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 power_grid_model_ds/_core/utils/misc.py,sha256=aAdlXjyKN6_T4b_CQZ879TbkbDZW4TwBUsxTpkwvymU,1228
 power_grid_model_ds/_core/utils/pickle.py,sha256=LGeTc7nu9RY1noOzLJzYaSHSWIgqzHy2xhmueKGuipc,1445
 power_grid_model_ds/_core/utils/zip.py,sha256=9RtJYhjlgNZOtxr4iI-CpsjT1axw5kCCqprfTjaIsiI,2197
-power_grid_model_ds-1.0.13.dist-info/LICENSE,sha256=GpbnG1pNl-ORtSP6dmHeiZvwy36UFli6MEZy1XlmBqo,14906
-power_grid_model_ds-1.0.13.dist-info/METADATA,sha256=G_AdWIdAk5jp2QL4pfinDGReAda1MeU-Y2gYosL8HeY,4270
-power_grid_model_ds-1.0.13.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-power_grid_model_ds-1.0.13.dist-info/top_level.txt,sha256=nJa103Eqvm5TESYEKPFVImfLg_ugGOBznikrM-rZQZg,20
-power_grid_model_ds-1.0.13.dist-info/RECORD,,
+power_grid_model_ds-1.1.1.dist-info/LICENSE,sha256=GpbnG1pNl-ORtSP6dmHeiZvwy36UFli6MEZy1XlmBqo,14906
+power_grid_model_ds-1.1.1.dist-info/METADATA,sha256=cby_jnH_ZVpRHkUu56M1NdbHpF7i5GyE2DQK0hj2p1k,4269
+power_grid_model_ds-1.1.1.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+power_grid_model_ds-1.1.1.dist-info/top_level.txt,sha256=nJa103Eqvm5TESYEKPFVImfLg_ugGOBznikrM-rZQZg,20
+power_grid_model_ds-1.1.1.dist-info/RECORD,,