From cbcad425fc7363f61387e9fa1ae42875db822d52 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Thu, 23 Jan 2025 20:06:53 +0100
Subject: [PATCH 1/5] feat: improve downstream nodes performance with local
 search

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 .../_core/model/graphs/models/base.py         |  21 ++++
 .../_core/model/graphs/models/rustworkx.py    |  11 ++
 .../_core/model/grids/base.py                 |  12 +--
 tests/performance/_constants.py               |  18 ++--
 tests/performance/_helpers.py                 | 101 +++++-------------
 tests/performance/array_performance_tests.py  |  54 +++++-----
 tests/performance/filter_performance_tests.py |  14 +--
 tests/performance/graph_performance_tests.py  |  19 ++--
 tests/performance/grid_performance_tests.py   |  30 ++++++
 tests/unit/model/grids/test_grid_search.py    |   3 +
 10 files changed, 156 insertions(+), 127 deletions(-)
 create mode 100644 tests/performance/grid_performance_tests.py

diff --git a/src/power_grid_model_ds/_core/model/graphs/models/base.py b/src/power_grid_model_ds/_core/model/graphs/models/base.py
index 2a4a577..0b817df 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/base.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/base.py
@@ -235,8 +235,26 @@ def get_connected(
             nodes_to_ignore=self._externals_to_internals(nodes_to_ignore),
             inclusive=inclusive,
         )
+
         return self._internals_to_externals(nodes)
 
+    def get_downstream_nodes(self, node_id: int, stop_node_ids: list[int], inclusive: bool = False) -> list[int]:
+        """Find all nodes connected to the node_id
+        args:
+            node_id: node id to start the search from
+            stop_node_ids: list of node ids to stop the search at
+            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
+        """
+        downstream_nodes = self._get_downstream_nodes(
+            node_id=self.external_to_internal(node_id),
+            stop_node_ids=self._externals_to_internals(stop_node_ids),
+            inclusive=inclusive,
+        )
+
+        return self._internals_to_externals(downstream_nodes)
+
     def find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph.
         Returns:
@@ -273,6 +291,9 @@ def _branch_is_relevant(self, branch: BranchArray) -> bool:
     @abstractmethod
     def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bool = False) -> list[int]: ...
 
+    @abstractmethod
+    def _get_downstream_nodes(self, node_id: int, stop_node_ids: list[int], inclusive: bool = False) -> list[int]: ...
+
     @abstractmethod
     def _has_branch(self, from_node_id, to_node_id) -> bool: ...
 
diff --git a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
index 6e1b7cd..2406bc2 100644
--- a/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
+++ b/src/power_grid_model_ds/_core/model/graphs/models/rustworkx.py
@@ -99,6 +99,15 @@ def _get_connected(self, node_id: int, nodes_to_ignore: list[int], inclusive: bo
 
         return connected_nodes
 
+    def _get_downstream_nodes(self, node_id: int, stop_node_ids: list[int], inclusive: bool = False) -> list[int]:
+        visitor = _NodeVisitor(stop_node_ids)
+        rx.bfs_search(self._graph, [node_id], visitor)
+        connected_nodes = visitor.nodes
+        path_to_substation, _ = self._get_shortest_path(node_id, visitor.discovered_nodes_to_ignore[0])
+        if inclusive:
+            _ = path_to_substation.pop(0)
+        return [node for node in connected_nodes if node not in path_to_substation]
+
     def _find_fundamental_cycles(self) -> list[list[int]]:
         """Find all fundamental cycles in the graph using Rustworkx.
 
@@ -112,8 +121,10 @@ class _NodeVisitor(BFSVisitor):
     def __init__(self, nodes_to_ignore: list[int]):
         self.nodes_to_ignore = nodes_to_ignore
         self.nodes: list[int] = []
+        self.discovered_nodes_to_ignore: list[int] = []
 
     def discover_vertex(self, v):
         if v in self.nodes_to_ignore:
+            self.discovered_nodes_to_ignore.append(v)
             raise PruneSearch
         self.nodes.append(v)
diff --git a/src/power_grid_model_ds/_core/model/grids/base.py b/src/power_grid_model_ds/_core/model/grids/base.py
index 0183ca1..18c94c8 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -331,6 +331,7 @@ 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
@@ -349,15 +350,14 @@ def get_downstream_nodes(self, node_id: int, inclusive: bool = False):
         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, stop_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
diff --git a/tests/performance/_constants.py b/tests/performance/_constants.py
index f114754..e531976 100644
--- a/tests/performance/_constants.py
+++ b/tests/performance/_constants.py
@@ -6,18 +6,18 @@
     "dtype = [('id', '<i8'), ('test_int', '<i8'), ('test_float', '<f8'), ('test_str', '<U50'), ('test_bool', '?')]; "
 )
 
-SETUP_CODES = {
-    "structured": "import numpy as np;" + NUMPY_DTYPE + "input_array = np.zeros({array_size}, dtype=dtype)",
-    "rec": "import numpy as np;" + NUMPY_DTYPE + "input_array = np.recarray(({array_size},),dtype=dtype)",
-    "fancy": "from tests.conftest import FancyTestArray; input_array=FancyTestArray.zeros({array_size});"
-    + "import numpy as np;input_array.id = np.arange({array_size})",
+ARRAY_SETUP_CODES = {
+    "structured": "import numpy as np;" + NUMPY_DTYPE + "input_array = np.zeros({size}, dtype=dtype)",
+    "rec": "import numpy as np;" + NUMPY_DTYPE + "input_array = np.recarray(({size},),dtype=dtype)",
+    "fancy": "from tests.conftest import FancyTestArray; input_array=FancyTestArray.zeros({size});"
+    + "import numpy as np;input_array.id = np.arange({size})",
 }
 
 GRAPH_SETUP_CODES = {
-    "rustworkx": "from power_grid_model_ds.model.grids.base import Grid;"
-    + "from power_grid_model_ds.data_source.generator.grid_generators import RadialGridGenerator;"
-    + "from power_grid_model_ds.model.graphs.models import RustworkxGraphModel;"
-    + "grid=RadialGridGenerator(nr_nodes={graph_size}, grid_class=Grid, graph_model=RustworkxGraphModel).run()",
+    "rustworkx": "from power_grid_model_ds import Grid;"
+    + "from power_grid_model_ds.generators import RadialGridGenerator;"
+    + "from power_grid_model_ds.graph_models import RustworkxGraphModel;"
+    + "grid=RadialGridGenerator(nr_nodes={size}, grid_class=Grid, graph_model=RustworkxGraphModel).run()",
 }
 
 SINGLE_REPEATS = 1000
diff --git a/tests/performance/_helpers.py b/tests/performance/_helpers.py
index 4e4b324..b36d9b0 100644
--- a/tests/performance/_helpers.py
+++ b/tests/performance/_helpers.py
@@ -4,85 +4,42 @@
 
 import inspect
 import timeit
-from typing import Generator
+from itertools import product
+from typing import Generator, Union
 
-from tests.performance._constants import GRAPH_SETUP_CODES, SETUP_CODES
-
-
-def do_performance_test(code_to_test: str | dict[str, str], array_sizes: list[int], repeats: int):
-    """Run the performance test for the given code."""
 
+def do_performance_test(
+    code_to_test: Union[str, dict[str, str], list[str]],
+    size_list: list[int],
+    repeats: int,
+    setup_codes: dict[str, str],
+):
+    """Generalized performance test runner."""
     print(f"{'-' * 20} {inspect.stack()[1][3]} {'-' * 20}")
 
-    for array_size in array_sizes:
+    for size in size_list:
+        formatted_setup_codes = {key: code.format(size=size) for key, code in setup_codes.items()}
         if isinstance(code_to_test, dict):
-            code_to_test_list = [code_to_test[variant].format(array_size=array_size) for variant in SETUP_CODES]
-        else:
-            code_to_test_list = [code_to_test.format(array_size=array_size)] * len(SETUP_CODES)
-        print(f"\n\tArray size: {array_size}\n")
-        setup_codes = [setup_code.format(array_size=array_size) for setup_code in SETUP_CODES.values()]
-        timings = _get_timings(setup_codes, code_to_test_list, repeats)
-
-        if code_to_test == "pass":
-            _print_timings(timings, list(SETUP_CODES.keys()), setup_codes)
+            code_to_test_list = [code_to_test[variant].format(size=size) for variant in setup_codes]
+            test_generator = zip(formatted_setup_codes.items(), code_to_test_list)
+        elif isinstance(code_to_test, list):
+            code_to_test_list = [code.format(size=size) for code in code_to_test]
+            test_generator = product(formatted_setup_codes.items(), code_to_test_list)
         else:
-            _print_timings(timings, list(SETUP_CODES.keys()), code_to_test_list)
-    print()
+            test_generator = product(formatted_setup_codes.items(), [code_to_test.format(size=size)])
 
+        print(f"\n\tsize: {size}\n")
 
-def do_graph_test(code_to_test: str | dict[str, str], graph_sizes: list[int], repeats: int):
-    """Run the performance test for the given code."""
+        timings = _get_timings(test_generator, repeats=repeats)
+        _print_timings(timings)
 
-    print(f"{'-' * 20} {inspect.stack()[1][3]} {'-' * 20}")
-
-    for graph_size in graph_sizes:
-        if isinstance(code_to_test, dict):
-            code_to_test_list = [code_to_test[variant] for variant in GRAPH_SETUP_CODES]
-        else:
-            code_to_test_list = [code_to_test] * len(GRAPH_SETUP_CODES)
-        print(f"\n\tGraph size: {graph_size}\n")
-        setup_codes = [setup_code.format(graph_size=graph_size) for setup_code in GRAPH_SETUP_CODES.values()]
-        timings = _get_timings(setup_codes, code_to_test_list, repeats)
-
-        if code_to_test == "pass":
-            _print_graph_timings(timings, list(GRAPH_SETUP_CODES.keys()), setup_codes)
-        else:
-            _print_graph_timings(timings, list(GRAPH_SETUP_CODES.keys()), code_to_test_list)
     print()
 
 
-def _print_test_code(code: str | dict[str, str], repeats: int):
-    print(f"{'-' * 40}")
-    if isinstance(code, dict):
-        for variant, code_variant in code.items():
-            print(f"{variant}")
-            print(f"\t{code_variant} (x {repeats})")
-        return
-    print(f"{code} (x {repeats})")
-
-
-def _print_graph_timings(timings: Generator, graph_types: list[str], code_list: list[str]):
-    for graph_type, timing, code in zip(graph_types, timings, code_list):
-        if ";" in code:
-            code = code.split(";")[-1]
-
-        code = code.replace("\n", " ").replace("\t", " ")
-        code = f"{graph_type}: " + code
-
-        if isinstance(timing, Exception):
-            print(f"\t\t{code.ljust(100)} | Not supported")
-            continue
-        print(f"\t\t{code.ljust(100)} | {sum(timing):.2f}s")
-
-
-def _print_timings(timings: Generator, array_types: list[str], code_list: list[str]):
-    for array, timing, code in zip(array_types, timings, code_list):
-        if ";" in code:
-            code = code.split(";")[-1]
-
-        code = code.replace("\n", " ").replace("\t", " ")
-        array_name = f"{array}_array"
-        code = code.replace("input_array", array_name)
+def _print_timings(timings: Generator):
+    for key, code, timing in timings:
+        code = code.split(";")[-1].replace("\n", " ").replace("\t", " ")
+        code = f"{key}: {code}"
 
         if isinstance(timing, Exception):
             print(f"\t\t{code.ljust(100)} | Not supported")
@@ -90,14 +47,14 @@ def _print_timings(timings: Generator, array_types: list[str], code_list: list[s
         print(f"\t\t{code.ljust(100)} | {sum(timing):.2f}s")
 
 
-def _get_timings(setup_codes: list[str], test_codes: list[str], repeats: int):
+def _get_timings(test_generator, repeats: int):
     """Return a generator with the timings for each array type."""
-    for setup_code, test_code in zip(setup_codes, test_codes):
+    for (key, setup_code), test_code in test_generator:
         if test_code == "pass":
-            yield timeit.repeat(setup_code, number=repeats)
+            yield key, "intialise", timeit.repeat(setup_code, number=repeats)
         else:
             try:
-                yield timeit.repeat(test_code, setup_code, number=repeats)
+                yield key, test_code, timeit.repeat(test_code, setup_code, number=repeats)
             # pylint: disable=broad-exception-caught
             except Exception as error:  # noqa
-                yield error
+                yield key, test_code, error
diff --git a/tests/performance/array_performance_tests.py b/tests/performance/array_performance_tests.py
index bd791b3..90161b7 100644
--- a/tests/performance/array_performance_tests.py
+++ b/tests/performance/array_performance_tests.py
@@ -11,18 +11,24 @@
 
 import logging
 
-from tests.performance._constants import ARRAY_SIZES_LARGE, ARRAY_SIZES_SMALL, LOOP_REPEATS, SINGLE_REPEATS
+from tests.performance._constants import (
+    ARRAY_SETUP_CODES,
+    ARRAY_SIZES_LARGE,
+    ARRAY_SIZES_SMALL,
+    LOOP_REPEATS,
+    SINGLE_REPEATS,
+)
 from tests.performance._helpers import do_performance_test
 
 logging.basicConfig(level=logging.INFO)
 
 
 def perftest_initialize():
-    do_performance_test("pass", ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test("pass", ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_slice():
-    do_performance_test("input_array[0:10]", ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test("input_array[0:10]", ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_set_attr():
@@ -31,77 +37,77 @@ def perftest_set_attr():
         "rec": "input_array.id = 1",
         "fancy": "input_array.id = 1",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_set_field():
-    do_performance_test("input_array['id'] = 1", ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test("input_array['id'] = 1", ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_slice_1():
-    code_to_test = "for i in range({array_size}): input_array[i]"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    code_to_test = "for i in range({size}): input_array[i]"
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_data_slice_1():
     code_to_test = {
-        "structured": "for i in range({array_size}): input_array[i]",
-        "rec": "for i in range({array_size}): input_array[i]",
-        "fancy": "for i in range({array_size}): input_array.data[i]",
+        "structured": "for i in range({size}): input_array[i]",
+        "rec": "for i in range({size}): input_array[i]",
+        "fancy": "for i in range({size}): input_array.data[i]",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_slice():
-    code_to_test = "for i in range({array_size}): input_array[i:i+1]"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    code_to_test = "for i in range({size}): input_array[i:i+1]"
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_set_field():
-    code_to_test = "for i in range({array_size}): input_array['id'][i] = 1"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    code_to_test = "for i in range({size}): input_array['id'][i] = 1"
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_get_field():
     code_to_test = "for row in input_array: row['id']"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_data_get_field():
     code_to_test = "for row in input_array.data: row['id']"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_get_attr():
     code_to_test = "for row in input_array: row.id"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100, ARRAY_SETUP_CODES)
 
 
 def perftest_fancypy_concat():
     code_to_test = {
         "structured": "import numpy as np;np.concatenate([input_array, input_array])",
         "rec": "import numpy as np;np.concatenate([input_array, input_array])",
-        "fancy": "import power_grid_model_ds._core.fancypy as fp;fp.concatenate(input_array, input_array)",
+        "fancy": "import power_grid_model_ds.fancypy as fp;fp.concatenate(input_array, input_array)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, 100)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, 100, ARRAY_SETUP_CODES)
 
 
 def perftest_fancypy_unique():
     code_to_test = {
         "structured": "import numpy as np;np.unique(input_array)",
         "rec": "import numpy as np;np.unique(input_array)",
-        "fancy": "import power_grid_model_ds._core.fancypy as fp;fp.unique(input_array)",
+        "fancy": "import power_grid_model_ds.fancypy as fp;fp.unique(input_array)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100, ARRAY_SETUP_CODES)
 
 
 def perftest_fancypy_sort():
     code_to_test = {
         "structured": "import numpy as np;np.sort(input_array)",
         "rec": "import numpy as np;np.sort(input_array)",
-        "fancy": "import power_grid_model_ds._core.fancypy as fp;fp.sort(input_array)",
+        "fancy": "import power_grid_model_ds.fancypy as fp;fp.sort(input_array)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100, ARRAY_SETUP_CODES)
 
 
 if __name__ == "__main__":
diff --git a/tests/performance/filter_performance_tests.py b/tests/performance/filter_performance_tests.py
index 86523ca..dbfebc8 100644
--- a/tests/performance/filter_performance_tests.py
+++ b/tests/performance/filter_performance_tests.py
@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 
-from tests.performance._constants import ARRAY_SIZES_LARGE, SINGLE_REPEATS
+from tests.performance._constants import ARRAY_SETUP_CODES, ARRAY_SIZES_LARGE, SINGLE_REPEATS
 from tests.performance._helpers import do_performance_test
 
 # pylint: disable=missing-function-docstring
@@ -14,7 +14,7 @@ def perftest_get():
         "rec": "input_array[np.isin(input_array['id'], 99)]",
         "fancy": "try:\n\tinput_array.get(id=99)\nexcept:\n\tpass",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_filter():
@@ -23,16 +23,16 @@ def perftest_filter():
         "rec": "input_array[np.isin(input_array['id'], 99)]",
         "fancy": "input_array.filter(id=99)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_update_by_id():
     code_to_test = {
-        "structured": "input_array['test_float'][np.isin(input_array['id'], np.arange({array_size}))] = 42.0",
-        "rec": "input_array['test_float'][np.isin(input_array['id'], np.arange({array_size}))] = 42.0",
-        "fancy": "input_array.update_by_id(ids=np.arange({array_size}), test_float=42.0, allow_missing = False)",
+        "structured": "input_array['test_float'][np.isin(input_array['id'], np.arange({size}))] = 42.0",
+        "rec": "input_array['test_float'][np.isin(input_array['id'], np.arange({size}))] = 42.0",
+        "fancy": "input_array.update_by_id(ids=np.arange({size}), test_float=42.0, allow_missing = False)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 if __name__ == "__main__":
diff --git a/tests/performance/graph_performance_tests.py b/tests/performance/graph_performance_tests.py
index 8b242df..9601165 100644
--- a/tests/performance/graph_performance_tests.py
+++ b/tests/performance/graph_performance_tests.py
@@ -2,7 +2,8 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 
-from tests.performance._helpers import do_graph_test
+from tests.performance._constants import GRAPH_SETUP_CODES
+from tests.performance._helpers import do_performance_test
 
 # pylint: disable=missing-function-docstring
 
@@ -11,40 +12,40 @@
 
 
 def perftest_initialize():
-    do_graph_test("pass", [10, 100], 100)
+    do_performance_test("pass", [10, 100], 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_get_components():
     code_to_test = (
-        "from power_grid_model_ds._core.model.enums.nodes import NodeType;"
+        "from power_grid_model_ds.enums import NodeType;"
         + "feeder_node_ids=grid.node.filter(node_type=NodeType.SUBSTATION_NODE).id;"
         + "grid.graphs.active_graph.get_components(feeder_node_ids)"
     )
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_set_feeder_ids():
     code_to_test = "grid.set_feeder_ids()"
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_delete_node():
     code_to_test = "grid.delete_node(grid.node[0]);"
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_from_arrays():
     code_to_test = "grid.graphs.complete_graph.__class__.from_arrays(grid);"
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_add_node():
     code_to_test = (
-        "from power_grid_model_ds._core.model.arrays import NodeArray;"
+        "from power_grid_model_ds.arrays import NodeArray;"
         + "new_node = NodeArray.zeros(1);"
         + "grid.add_node(node=new_node)"
     )
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 if __name__ == "__main__":
diff --git a/tests/performance/grid_performance_tests.py b/tests/performance/grid_performance_tests.py
new file mode 100644
index 0000000..5d39f93
--- /dev/null
+++ b/tests/performance/grid_performance_tests.py
@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+from tests.performance._helpers import do_performance_test
+
+# pylint: disable=missing-function-docstring
+
+
+def test_get_downstream_nodes_performance():
+    setup_code = {
+        "grid": "import numpy as np;"
+        + "from power_grid_model_ds.enums import NodeType;"
+        + "from power_grid_model_ds import Grid;"
+        + "from power_grid_model_ds.generators import RadialGridGenerator;"
+        + "from power_grid_model_ds.graph_models import RustworkxGraphModel;"
+        + "grid=RadialGridGenerator(nr_nodes={size}, grid_class=Grid, graph_model=RustworkxGraphModel).run();"
+        + "non_substation_node = grid.arrays.nodes.filter(node_type=NodeType.UNSPECIFIED).id;"
+        + "node_id = np.random.choice(non_substation_node)"
+    }
+
+    code_to_test = [
+        "grid.get_downstream_nodes(node_id)",
+    ]
+
+    do_performance_test(code_to_test, [10, 1000, 5000], 100, setup_code)
+
+
+if __name__ == "__main__":
+    test_get_downstream_nodes_performance()
diff --git a/tests/unit/model/grids/test_grid_search.py b/tests/unit/model/grids/test_grid_search.py
index f3ac243..1ed1e06 100644
--- a/tests/unit/model/grids/test_grid_search.py
+++ b/tests/unit/model/grids/test_grid_search.py
@@ -30,6 +30,9 @@ def test_get_downstream_nodes(basic_grid):
     downstream_nodes = basic_grid.get_downstream_nodes(node_id=102)
     assert {103, 106} == set(downstream_nodes)
 
+    downstream_nodes = basic_grid.get_downstream_nodes(node_id=102, inclusive=True)
+    assert {102, 103, 106} == set(downstream_nodes)
+
 
 def test_get_downstream_nodes_from_substation_node(basic_grid):
     """Test that get_downstream_nodes raises the expected error when

From 3d72bab35f8b639c5c3ba02cb079accce2ac4d08 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Fri, 24 Jan 2025 13:17:43 +0100
Subject: [PATCH 2/5] test: add testing on sorting

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 tests/unit/model/grids/test_grid_search.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/tests/unit/model/grids/test_grid_search.py b/tests/unit/model/grids/test_grid_search.py
index 1ed1e06..034f1cd 100644
--- a/tests/unit/model/grids/test_grid_search.py
+++ b/tests/unit/model/grids/test_grid_search.py
@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 
+from power_grid_model_ds import Grid
 from power_grid_model_ds._core.model.arrays.base.errors import RecordDoesNotExist
 from power_grid_model_ds._core.model.enums.nodes import NodeType
 
@@ -25,13 +26,19 @@ def test_grid_get_nearest_substation_node_no_substation(basic_grid):
         basic_grid.get_nearest_substation_node(node_id=103)
 
 
-def test_get_downstream_nodes(basic_grid):
+def test_get_downstream_nodes(basic_grid: Grid):
     """Test that get_downstream_nodes returns the expected nodes."""
+    # Move the open line to be able to test sorting of nodes by distance correctly
+    basic_grid.make_active(basic_grid.line.get(203))
+    basic_grid.make_inactive(basic_grid.link.get(601))
     downstream_nodes = basic_grid.get_downstream_nodes(node_id=102)
-    assert {103, 106} == set(downstream_nodes)
+    assert downstream_nodes[-1] == 104  # Furthest away
+    assert {103, 104, 106} == set(downstream_nodes)
 
     downstream_nodes = basic_grid.get_downstream_nodes(node_id=102, inclusive=True)
-    assert {102, 103, 106} == set(downstream_nodes)
+    assert downstream_nodes[0] == 102
+    assert downstream_nodes[-1] == 104
+    assert {102, 103, 104, 106} == set(downstream_nodes)
 
 
 def test_get_downstream_nodes_from_substation_node(basic_grid):

From da0c33d1a39a5897819a3465d26b834c6e33a113 Mon Sep 17 00:00:00 2001
From: jaapschoutenalliander <jaap.schouten@alliander.com>
Date: Tue, 28 Jan 2025 19:33:34 +0100
Subject: [PATCH 3/5] chore: update performance test

Signed-off-by: jaapschoutenalliander <jaap.schouten@alliander.com>
---
 tests/performance/grid_performance_tests.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/performance/grid_performance_tests.py b/tests/performance/grid_performance_tests.py
index 5d39f93..db2fc28 100644
--- a/tests/performance/grid_performance_tests.py
+++ b/tests/performance/grid_performance_tests.py
@@ -15,7 +15,7 @@ def test_get_downstream_nodes_performance():
         + "from power_grid_model_ds.generators import RadialGridGenerator;"
         + "from power_grid_model_ds.graph_models import RustworkxGraphModel;"
         + "grid=RadialGridGenerator(nr_nodes={size}, grid_class=Grid, graph_model=RustworkxGraphModel).run();"
-        + "non_substation_node = grid.arrays.nodes.filter(node_type=NodeType.UNSPECIFIED).id;"
+        + "non_substation_node = grid.node.filter(node_type=NodeType.UNSPECIFIED).id;"
         + "node_id = np.random.choice(non_substation_node)"
     }
 

From fc99bc0e5c2d1b4e9d1119a69ce6b2e64d5e37a8 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Wed, 29 Jan 2025 11:09:15 +0100
Subject: [PATCH 4/5] Update test_get_downstream_nodes

Co-authored-by: jaapschoutenalliander <jaap.schouten@alliander.com>
Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 tests/unit/model/grids/test_grid_search.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/tests/unit/model/grids/test_grid_search.py b/tests/unit/model/grids/test_grid_search.py
index f3ac243..034f1cd 100644
--- a/tests/unit/model/grids/test_grid_search.py
+++ b/tests/unit/model/grids/test_grid_search.py
@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 
+from power_grid_model_ds import Grid
 from power_grid_model_ds._core.model.arrays.base.errors import RecordDoesNotExist
 from power_grid_model_ds._core.model.enums.nodes import NodeType
 
@@ -25,10 +26,19 @@ def test_grid_get_nearest_substation_node_no_substation(basic_grid):
         basic_grid.get_nearest_substation_node(node_id=103)
 
 
-def test_get_downstream_nodes(basic_grid):
+def test_get_downstream_nodes(basic_grid: Grid):
     """Test that get_downstream_nodes returns the expected nodes."""
+    # Move the open line to be able to test sorting of nodes by distance correctly
+    basic_grid.make_active(basic_grid.line.get(203))
+    basic_grid.make_inactive(basic_grid.link.get(601))
     downstream_nodes = basic_grid.get_downstream_nodes(node_id=102)
-    assert {103, 106} == set(downstream_nodes)
+    assert downstream_nodes[-1] == 104  # Furthest away
+    assert {103, 104, 106} == set(downstream_nodes)
+
+    downstream_nodes = basic_grid.get_downstream_nodes(node_id=102, inclusive=True)
+    assert downstream_nodes[0] == 102
+    assert downstream_nodes[-1] == 104
+    assert {102, 103, 104, 106} == set(downstream_nodes)
 
 
 def test_get_downstream_nodes_from_substation_node(basic_grid):

From 7cb6fa24a5c6fd5790783dcace41b8e76cf6b2e6 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Wed, 29 Jan 2025 11:09:24 +0100
Subject: [PATCH 5/5] Update performance tests

Co-authored-by: jaapschoutenalliander <jaap.schouten@alliander.com>
Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 tests/performance/_constants.py               |  18 ++--
 tests/performance/_helpers.py                 | 101 +++++-------------
 tests/performance/array_performance_tests.py  |  54 +++++-----
 tests/performance/filter_performance_tests.py |  14 +--
 tests/performance/graph_performance_tests.py  |  19 ++--
 tests/performance/grid_performance_tests.py   |  30 ++++++
 6 files changed, 115 insertions(+), 121 deletions(-)
 create mode 100644 tests/performance/grid_performance_tests.py

diff --git a/tests/performance/_constants.py b/tests/performance/_constants.py
index f114754..e531976 100644
--- a/tests/performance/_constants.py
+++ b/tests/performance/_constants.py
@@ -6,18 +6,18 @@
     "dtype = [('id', '<i8'), ('test_int', '<i8'), ('test_float', '<f8'), ('test_str', '<U50'), ('test_bool', '?')]; "
 )
 
-SETUP_CODES = {
-    "structured": "import numpy as np;" + NUMPY_DTYPE + "input_array = np.zeros({array_size}, dtype=dtype)",
-    "rec": "import numpy as np;" + NUMPY_DTYPE + "input_array = np.recarray(({array_size},),dtype=dtype)",
-    "fancy": "from tests.conftest import FancyTestArray; input_array=FancyTestArray.zeros({array_size});"
-    + "import numpy as np;input_array.id = np.arange({array_size})",
+ARRAY_SETUP_CODES = {
+    "structured": "import numpy as np;" + NUMPY_DTYPE + "input_array = np.zeros({size}, dtype=dtype)",
+    "rec": "import numpy as np;" + NUMPY_DTYPE + "input_array = np.recarray(({size},),dtype=dtype)",
+    "fancy": "from tests.conftest import FancyTestArray; input_array=FancyTestArray.zeros({size});"
+    + "import numpy as np;input_array.id = np.arange({size})",
 }
 
 GRAPH_SETUP_CODES = {
-    "rustworkx": "from power_grid_model_ds.model.grids.base import Grid;"
-    + "from power_grid_model_ds.data_source.generator.grid_generators import RadialGridGenerator;"
-    + "from power_grid_model_ds.model.graphs.models import RustworkxGraphModel;"
-    + "grid=RadialGridGenerator(nr_nodes={graph_size}, grid_class=Grid, graph_model=RustworkxGraphModel).run()",
+    "rustworkx": "from power_grid_model_ds import Grid;"
+    + "from power_grid_model_ds.generators import RadialGridGenerator;"
+    + "from power_grid_model_ds.graph_models import RustworkxGraphModel;"
+    + "grid=RadialGridGenerator(nr_nodes={size}, grid_class=Grid, graph_model=RustworkxGraphModel).run()",
 }
 
 SINGLE_REPEATS = 1000
diff --git a/tests/performance/_helpers.py b/tests/performance/_helpers.py
index 4e4b324..b36d9b0 100644
--- a/tests/performance/_helpers.py
+++ b/tests/performance/_helpers.py
@@ -4,85 +4,42 @@
 
 import inspect
 import timeit
-from typing import Generator
+from itertools import product
+from typing import Generator, Union
 
-from tests.performance._constants import GRAPH_SETUP_CODES, SETUP_CODES
-
-
-def do_performance_test(code_to_test: str | dict[str, str], array_sizes: list[int], repeats: int):
-    """Run the performance test for the given code."""
 
+def do_performance_test(
+    code_to_test: Union[str, dict[str, str], list[str]],
+    size_list: list[int],
+    repeats: int,
+    setup_codes: dict[str, str],
+):
+    """Generalized performance test runner."""
     print(f"{'-' * 20} {inspect.stack()[1][3]} {'-' * 20}")
 
-    for array_size in array_sizes:
+    for size in size_list:
+        formatted_setup_codes = {key: code.format(size=size) for key, code in setup_codes.items()}
         if isinstance(code_to_test, dict):
-            code_to_test_list = [code_to_test[variant].format(array_size=array_size) for variant in SETUP_CODES]
-        else:
-            code_to_test_list = [code_to_test.format(array_size=array_size)] * len(SETUP_CODES)
-        print(f"\n\tArray size: {array_size}\n")
-        setup_codes = [setup_code.format(array_size=array_size) for setup_code in SETUP_CODES.values()]
-        timings = _get_timings(setup_codes, code_to_test_list, repeats)
-
-        if code_to_test == "pass":
-            _print_timings(timings, list(SETUP_CODES.keys()), setup_codes)
+            code_to_test_list = [code_to_test[variant].format(size=size) for variant in setup_codes]
+            test_generator = zip(formatted_setup_codes.items(), code_to_test_list)
+        elif isinstance(code_to_test, list):
+            code_to_test_list = [code.format(size=size) for code in code_to_test]
+            test_generator = product(formatted_setup_codes.items(), code_to_test_list)
         else:
-            _print_timings(timings, list(SETUP_CODES.keys()), code_to_test_list)
-    print()
+            test_generator = product(formatted_setup_codes.items(), [code_to_test.format(size=size)])
 
+        print(f"\n\tsize: {size}\n")
 
-def do_graph_test(code_to_test: str | dict[str, str], graph_sizes: list[int], repeats: int):
-    """Run the performance test for the given code."""
+        timings = _get_timings(test_generator, repeats=repeats)
+        _print_timings(timings)
 
-    print(f"{'-' * 20} {inspect.stack()[1][3]} {'-' * 20}")
-
-    for graph_size in graph_sizes:
-        if isinstance(code_to_test, dict):
-            code_to_test_list = [code_to_test[variant] for variant in GRAPH_SETUP_CODES]
-        else:
-            code_to_test_list = [code_to_test] * len(GRAPH_SETUP_CODES)
-        print(f"\n\tGraph size: {graph_size}\n")
-        setup_codes = [setup_code.format(graph_size=graph_size) for setup_code in GRAPH_SETUP_CODES.values()]
-        timings = _get_timings(setup_codes, code_to_test_list, repeats)
-
-        if code_to_test == "pass":
-            _print_graph_timings(timings, list(GRAPH_SETUP_CODES.keys()), setup_codes)
-        else:
-            _print_graph_timings(timings, list(GRAPH_SETUP_CODES.keys()), code_to_test_list)
     print()
 
 
-def _print_test_code(code: str | dict[str, str], repeats: int):
-    print(f"{'-' * 40}")
-    if isinstance(code, dict):
-        for variant, code_variant in code.items():
-            print(f"{variant}")
-            print(f"\t{code_variant} (x {repeats})")
-        return
-    print(f"{code} (x {repeats})")
-
-
-def _print_graph_timings(timings: Generator, graph_types: list[str], code_list: list[str]):
-    for graph_type, timing, code in zip(graph_types, timings, code_list):
-        if ";" in code:
-            code = code.split(";")[-1]
-
-        code = code.replace("\n", " ").replace("\t", " ")
-        code = f"{graph_type}: " + code
-
-        if isinstance(timing, Exception):
-            print(f"\t\t{code.ljust(100)} | Not supported")
-            continue
-        print(f"\t\t{code.ljust(100)} | {sum(timing):.2f}s")
-
-
-def _print_timings(timings: Generator, array_types: list[str], code_list: list[str]):
-    for array, timing, code in zip(array_types, timings, code_list):
-        if ";" in code:
-            code = code.split(";")[-1]
-
-        code = code.replace("\n", " ").replace("\t", " ")
-        array_name = f"{array}_array"
-        code = code.replace("input_array", array_name)
+def _print_timings(timings: Generator):
+    for key, code, timing in timings:
+        code = code.split(";")[-1].replace("\n", " ").replace("\t", " ")
+        code = f"{key}: {code}"
 
         if isinstance(timing, Exception):
             print(f"\t\t{code.ljust(100)} | Not supported")
@@ -90,14 +47,14 @@ def _print_timings(timings: Generator, array_types: list[str], code_list: list[s
         print(f"\t\t{code.ljust(100)} | {sum(timing):.2f}s")
 
 
-def _get_timings(setup_codes: list[str], test_codes: list[str], repeats: int):
+def _get_timings(test_generator, repeats: int):
     """Return a generator with the timings for each array type."""
-    for setup_code, test_code in zip(setup_codes, test_codes):
+    for (key, setup_code), test_code in test_generator:
         if test_code == "pass":
-            yield timeit.repeat(setup_code, number=repeats)
+            yield key, "intialise", timeit.repeat(setup_code, number=repeats)
         else:
             try:
-                yield timeit.repeat(test_code, setup_code, number=repeats)
+                yield key, test_code, timeit.repeat(test_code, setup_code, number=repeats)
             # pylint: disable=broad-exception-caught
             except Exception as error:  # noqa
-                yield error
+                yield key, test_code, error
diff --git a/tests/performance/array_performance_tests.py b/tests/performance/array_performance_tests.py
index bd791b3..90161b7 100644
--- a/tests/performance/array_performance_tests.py
+++ b/tests/performance/array_performance_tests.py
@@ -11,18 +11,24 @@
 
 import logging
 
-from tests.performance._constants import ARRAY_SIZES_LARGE, ARRAY_SIZES_SMALL, LOOP_REPEATS, SINGLE_REPEATS
+from tests.performance._constants import (
+    ARRAY_SETUP_CODES,
+    ARRAY_SIZES_LARGE,
+    ARRAY_SIZES_SMALL,
+    LOOP_REPEATS,
+    SINGLE_REPEATS,
+)
 from tests.performance._helpers import do_performance_test
 
 logging.basicConfig(level=logging.INFO)
 
 
 def perftest_initialize():
-    do_performance_test("pass", ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test("pass", ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_slice():
-    do_performance_test("input_array[0:10]", ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test("input_array[0:10]", ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_set_attr():
@@ -31,77 +37,77 @@ def perftest_set_attr():
         "rec": "input_array.id = 1",
         "fancy": "input_array.id = 1",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_set_field():
-    do_performance_test("input_array['id'] = 1", ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test("input_array['id'] = 1", ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_slice_1():
-    code_to_test = "for i in range({array_size}): input_array[i]"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    code_to_test = "for i in range({size}): input_array[i]"
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_data_slice_1():
     code_to_test = {
-        "structured": "for i in range({array_size}): input_array[i]",
-        "rec": "for i in range({array_size}): input_array[i]",
-        "fancy": "for i in range({array_size}): input_array.data[i]",
+        "structured": "for i in range({size}): input_array[i]",
+        "rec": "for i in range({size}): input_array[i]",
+        "fancy": "for i in range({size}): input_array.data[i]",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_slice():
-    code_to_test = "for i in range({array_size}): input_array[i:i+1]"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    code_to_test = "for i in range({size}): input_array[i:i+1]"
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_set_field():
-    code_to_test = "for i in range({array_size}): input_array['id'][i] = 1"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    code_to_test = "for i in range({size}): input_array['id'][i] = 1"
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_get_field():
     code_to_test = "for row in input_array: row['id']"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_data_get_field():
     code_to_test = "for row in input_array.data: row['id']"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, LOOP_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_loop_get_attr():
     code_to_test = "for row in input_array: row.id"
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100, ARRAY_SETUP_CODES)
 
 
 def perftest_fancypy_concat():
     code_to_test = {
         "structured": "import numpy as np;np.concatenate([input_array, input_array])",
         "rec": "import numpy as np;np.concatenate([input_array, input_array])",
-        "fancy": "import power_grid_model_ds._core.fancypy as fp;fp.concatenate(input_array, input_array)",
+        "fancy": "import power_grid_model_ds.fancypy as fp;fp.concatenate(input_array, input_array)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, 100)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, 100, ARRAY_SETUP_CODES)
 
 
 def perftest_fancypy_unique():
     code_to_test = {
         "structured": "import numpy as np;np.unique(input_array)",
         "rec": "import numpy as np;np.unique(input_array)",
-        "fancy": "import power_grid_model_ds._core.fancypy as fp;fp.unique(input_array)",
+        "fancy": "import power_grid_model_ds.fancypy as fp;fp.unique(input_array)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100, ARRAY_SETUP_CODES)
 
 
 def perftest_fancypy_sort():
     code_to_test = {
         "structured": "import numpy as np;np.sort(input_array)",
         "rec": "import numpy as np;np.sort(input_array)",
-        "fancy": "import power_grid_model_ds._core.fancypy as fp;fp.sort(input_array)",
+        "fancy": "import power_grid_model_ds.fancypy as fp;fp.sort(input_array)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100)
+    do_performance_test(code_to_test, ARRAY_SIZES_SMALL, 100, ARRAY_SETUP_CODES)
 
 
 if __name__ == "__main__":
diff --git a/tests/performance/filter_performance_tests.py b/tests/performance/filter_performance_tests.py
index 86523ca..dbfebc8 100644
--- a/tests/performance/filter_performance_tests.py
+++ b/tests/performance/filter_performance_tests.py
@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 
-from tests.performance._constants import ARRAY_SIZES_LARGE, SINGLE_REPEATS
+from tests.performance._constants import ARRAY_SETUP_CODES, ARRAY_SIZES_LARGE, SINGLE_REPEATS
 from tests.performance._helpers import do_performance_test
 
 # pylint: disable=missing-function-docstring
@@ -14,7 +14,7 @@ def perftest_get():
         "rec": "input_array[np.isin(input_array['id'], 99)]",
         "fancy": "try:\n\tinput_array.get(id=99)\nexcept:\n\tpass",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_filter():
@@ -23,16 +23,16 @@ def perftest_filter():
         "rec": "input_array[np.isin(input_array['id'], 99)]",
         "fancy": "input_array.filter(id=99)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 def perftest_update_by_id():
     code_to_test = {
-        "structured": "input_array['test_float'][np.isin(input_array['id'], np.arange({array_size}))] = 42.0",
-        "rec": "input_array['test_float'][np.isin(input_array['id'], np.arange({array_size}))] = 42.0",
-        "fancy": "input_array.update_by_id(ids=np.arange({array_size}), test_float=42.0, allow_missing = False)",
+        "structured": "input_array['test_float'][np.isin(input_array['id'], np.arange({size}))] = 42.0",
+        "rec": "input_array['test_float'][np.isin(input_array['id'], np.arange({size}))] = 42.0",
+        "fancy": "input_array.update_by_id(ids=np.arange({size}), test_float=42.0, allow_missing = False)",
     }
-    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS)
+    do_performance_test(code_to_test, ARRAY_SIZES_LARGE, SINGLE_REPEATS, ARRAY_SETUP_CODES)
 
 
 if __name__ == "__main__":
diff --git a/tests/performance/graph_performance_tests.py b/tests/performance/graph_performance_tests.py
index 8b242df..9601165 100644
--- a/tests/performance/graph_performance_tests.py
+++ b/tests/performance/graph_performance_tests.py
@@ -2,7 +2,8 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 
-from tests.performance._helpers import do_graph_test
+from tests.performance._constants import GRAPH_SETUP_CODES
+from tests.performance._helpers import do_performance_test
 
 # pylint: disable=missing-function-docstring
 
@@ -11,40 +12,40 @@
 
 
 def perftest_initialize():
-    do_graph_test("pass", [10, 100], 100)
+    do_performance_test("pass", [10, 100], 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_get_components():
     code_to_test = (
-        "from power_grid_model_ds._core.model.enums.nodes import NodeType;"
+        "from power_grid_model_ds.enums import NodeType;"
         + "feeder_node_ids=grid.node.filter(node_type=NodeType.SUBSTATION_NODE).id;"
         + "grid.graphs.active_graph.get_components(feeder_node_ids)"
     )
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_set_feeder_ids():
     code_to_test = "grid.set_feeder_ids()"
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_delete_node():
     code_to_test = "grid.delete_node(grid.node[0]);"
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_from_arrays():
     code_to_test = "grid.graphs.complete_graph.__class__.from_arrays(grid);"
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 def perftest_add_node():
     code_to_test = (
-        "from power_grid_model_ds._core.model.arrays import NodeArray;"
+        "from power_grid_model_ds.arrays import NodeArray;"
         + "new_node = NodeArray.zeros(1);"
         + "grid.add_node(node=new_node)"
     )
-    do_graph_test(code_to_test, GRAPH_SIZES, 100)
+    do_performance_test(code_to_test, GRAPH_SIZES, 100, setup_codes=GRAPH_SETUP_CODES)
 
 
 if __name__ == "__main__":
diff --git a/tests/performance/grid_performance_tests.py b/tests/performance/grid_performance_tests.py
new file mode 100644
index 0000000..db2fc28
--- /dev/null
+++ b/tests/performance/grid_performance_tests.py
@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <powergridmodel@lfenergy.org>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+from tests.performance._helpers import do_performance_test
+
+# pylint: disable=missing-function-docstring
+
+
+def test_get_downstream_nodes_performance():
+    setup_code = {
+        "grid": "import numpy as np;"
+        + "from power_grid_model_ds.enums import NodeType;"
+        + "from power_grid_model_ds import Grid;"
+        + "from power_grid_model_ds.generators import RadialGridGenerator;"
+        + "from power_grid_model_ds.graph_models import RustworkxGraphModel;"
+        + "grid=RadialGridGenerator(nr_nodes={size}, grid_class=Grid, graph_model=RustworkxGraphModel).run();"
+        + "non_substation_node = grid.node.filter(node_type=NodeType.UNSPECIFIED).id;"
+        + "node_id = np.random.choice(non_substation_node)"
+    }
+
+    code_to_test = [
+        "grid.get_downstream_nodes(node_id)",
+    ]
+
+    do_performance_test(code_to_test, [10, 1000, 5000], 100, setup_code)
+
+
+if __name__ == "__main__":
+    test_get_downstream_nodes_performance()