merge

Signed-off-by: Thijs Baaijen <[email protected]>

Author: Thijss

src/power_grid_model_ds/_core/model/grids/_text_sources.py

@@ -5,6 +5,7 @@
 """Create a grid from text a text file"""
 
 import logging
+import re
 from typing import TYPE_CHECKING
 
 from power_grid_model_ds._core.model.enums.nodes import NodeType
@@ -34,9 +35,12 @@ class TextSource:
     def __init__(self, grid_class: type["Grid"]):
         self.grid = grid_class.empty()
 
-    def load_from_txt(self, txt_lines: list[str]) -> "Grid":
+    def load_from_txt(self, *args: str) -> "Grid":
         """Load a grid from text"""
-        txt_nodes, txt_branches = self.read_txt(txt_lines)
+
+        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()
@@ -50,11 +54,14 @@ def read_txt(txt_lines: list[str]) -> tuple[set, dict]:
         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(" ")
-            except ValueError as err:
-                raise ValueError(f"Text line '{text_line}' is invalid. Skipping...") from err
-            comments = comments[0].split(",") if comments else []
+
+            pattern = re.compile(r"^\s*(\S+)\s+(\S+)(?:\s+(\S+))?\s*$")
+            match = pattern.match(text_line)
+            if not match:
+                raise ValueError(f"Text line '{text_line}' is invalid. Skipping...")
+
+            from_node_str, to_node_str, comment = match.groups()
+            comments = comment.split(",") if comment else []
 
             txt_nodes |= {from_node_str, to_node_str}
             txt_branches[(from_node_str, to_node_str)] = comments

src/power_grid_model_ds/_core/model/grids/base.py

@@ -331,7 +331,6 @@ def get_nearest_substation_node(self, node_id: int):
 
     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
@@ -350,14 +349,15 @@ def get_downstream_nodes(self, node_id: int, inclusive: bool = False):
         Returns:
             list[int]: The downstream nodes.
         """
-        substation_nodes = self.node.filter(node_type=NodeType.SUBSTATION_NODE.value)
+        substation_node_id = self.get_nearest_substation_node(node_id).id.item()
 
-        if node_id in substation_nodes.id:
+        if node_id == substation_node_id:
             raise NotImplementedError("get_downstream_nodes is not implemented for substation nodes!")
 
-        return self.graphs.active_graph.get_downstream_nodes(
-            node_id=node_id, start_node_ids=list(substation_nodes.id), inclusive=inclusive
-        )
+        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)
 
     def cache(self, cache_dir: Path, cache_name: str, compress: bool = True):
         """Cache Grid to a folder
@@ -409,15 +409,19 @@ def _from_pickle(cls, pickle_path: Path):
         return grid
 
     @classmethod
-    def from_txt(cls, txt_lines: list[str]):
+    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
 
-        Example:
-            >>> Grid.from_txt(["1 2", "2 3", "3 4 transformer", "4 5", "S1 6"])
+        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(txt_lines)
+        return TextSource(grid_class=cls).load_from_txt(*args)
 
     @classmethod
     # pylint: disable=arguments-differ
@@ -429,7 +433,7 @@ def from_txt_file(cls, txt_file_path: 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)
+        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"""

tests/unit/model/grids/test_grid_base.py

@@ -277,7 +277,7 @@ def test_grid_as_str(basic_grid):
 
 class TestFromTxt:
     def test_from_txt_lines(self):
-        txt_lines = [
+        grid = Grid.from_txt(
             "S1 2",
             "S1 3 open",
             "2 7",
@@ -286,64 +286,55 @@ def test_from_txt_lines(self):
             "5 7",
             "7 8",
             "8 9",
-        ]
-        grid = Grid.from_txt(txt_lines)
+        )
         assert 8 == grid.node.size
         assert 1 == grid.branches.filter(to_status=0).size
         assert 1 == grid.transformer.size
         np.testing.assert_array_equal([14, 10, 11, 12, 13, 15, 16, 17], grid.branches.id)
 
+    def test_from_txt_string(self):
+        txt_string = "S1 2\nS1 3 open\n2 7\n3 5\n3 6 transformer\n5 7\n7 8\n8 9"
+        assert Grid.from_txt(txt_string)
+
+    def test_from_txt_string_with_spaces(self):
+        txt_string = "S1 2     \nS1 3   open\n2    7\n3 5\n   3 6 transformer\n5 7\n7   8\n8 9"
+        assert Grid.from_txt(txt_string)
+
+    def test_from_docstring(self):
+        assert Grid.from_txt("""
+        S1 2
+        S1 3 open
+        2 7
+        3 5
+        3 6 transformer
+        5 7
+        7 8
+        8 9
+        """)
+
     def test_from_txt_with_branch_ids(self):
-        txt_lines = [
-            "S1 2 91",
-            "S1 3 92,open",
-            "2 7 93",
-            "3 5 94",
-            "3 6 transformer,95",
-            "5 7 96",
-            "7 8 97",
-            "8 9 98",
-        ]
-
-        grid = Grid.from_txt(txt_lines)
+        grid = Grid.from_txt(
+            "S1 2 91", "S1 3 92,open", "2 7 93", "3 5 94", "3 6 transformer,95", "5 7 96", "7 8 97", "8 9 98"
+        )
         assert 8 == grid.node.size
         assert 1 == grid.branches.filter(to_status=0).size
         assert 1 == grid.transformer.size
         np.testing.assert_array_equal([95, 91, 92, 93, 94, 96, 97, 98], grid.branches.id)
 
     def test_from_txt_with_conflicting_ids(self):
-        txt_lines = [
-            "S1 2",
-            "1 3",
-        ]
+        with pytest.raises(ValueError):
+            Grid.from_txt("S1 2", "1 3")
 
+    def test_from_txt_with_invalid_line(self):
         with pytest.raises(ValueError):
-            Grid.from_txt(txt_lines)
+            Grid.from_txt("S1 2 arg3 arg4")
 
     def test_from_txt_with_unordered_node_ids(self):
-        txt_lines = [
-            "S1 2",
-            "S1 10",
-            "10 11",
-            "2 5",
-            "5 6",
-            "3 4",
-            "3 7",
-        ]
-        grid = Grid.from_txt(txt_lines)
+        grid = Grid.from_txt("S1 2", "S1 10", "10 11", "2 5", "5 6", "3 4", "3 7")
         assert 9 == grid.node.size
 
     def test_from_txt_with_unordered_branch_ids(self):
-        txt_lines = [
-            "5 6 16",
-            "3 4 17",
-            "3 7 18",
-            "S1 2 12",
-            "S1 10 13",
-            "10 11 14",
-            "2 5 15",
-        ]
-        grid = Grid.from_txt(txt_lines)
+        grid = Grid.from_txt("5 6 16", "3 4 17", "3 7 18", "S1 2 12", "S1 10 13", "10 11 14", "2 5 15")
         assert 9 == grid.node.size
 
     def test_from_txt_file(self, tmp_path):