feat: support all component types in interface (#114)

* feat: support all component types in interface

Signed-off-by: Vincent Koppen <vincent.koppen@alliander.com>

* linting

Signed-off-by: Vincent Koppen <vincent.koppen@alliander.com>

---------

Signed-off-by: Vincent Koppen <vincent.koppen@alliander.com>

Author: vincentkoppen

docs/demos/connecting_new_consumer.ipynb

@@ -131,7 +131,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from power_grid_model_ds._core.load_flow import PowerGridModelInterface\n",
+    "from power_grid_model_ds import PowerGridModelInterface\n",
     "\n",
     "R_PER_KM = 0.1\n",
     "X_PER_KM = 0.1\n",

src/power_grid_model_ds/__init__.py

@@ -2,8 +2,8 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 
-from power_grid_model_ds._core.load_flow import PowerGridModelInterface
 from power_grid_model_ds._core.model.graphs.container import GraphContainer
 from power_grid_model_ds._core.model.grids.base import Grid
+from power_grid_model_ds._core.power_grid_model_interface import PowerGridModelInterface
 
 __all__ = ["Grid", "GraphContainer", "PowerGridModelInterface"]

…c/power_grid_model_ds/_core/load_flow.py …l_ds/_core/power_grid_model_interface.pysrc/power_grid_model_ds/_core/load_flow.py renamed to src/power_grid_model_ds/_core/power_grid_model_interface.py src/power_grid_model_ds/_core/load_flow.py renamed to src/power_grid_model_ds/_core/power_grid_model_interface.py

@@ -9,25 +9,10 @@
 
 import numpy as np
 from numpy.typing import NDArray
-from power_grid_model import CalculationMethod, PowerGridModel, initialize_array
+from power_grid_model import CalculationMethod, ComponentType, PowerGridModel, initialize_array
 
 from power_grid_model_ds._core.model.grids.base import Grid
 
-PGM_ARRAYS = [
-    "node",
-    "line",
-    "link",
-    "transformer",
-    "three_winding_transformer",
-    "sym_load",
-    "sym_gen",
-    "source",
-    "transformer_tap_regulator",
-    "sym_power_sensor",
-    "sym_voltage_sensor",
-    "asym_voltage_sensor",
-]
-
 
 class PGMCoreException(Exception):
     """Raised when there is an error in running the power grid model"""
@@ -70,7 +55,9 @@ def create_input_from_grid(self):
         """
         Create input for the PowerGridModel
         """
-        for array_name in PGM_ARRAYS:
+        for array_name in ComponentType:
+            if not hasattr(self.grid, array_name):
+                continue
             pgm_array = self._create_power_grid_array(array_name=array_name)
             self._input_data[array_name] = pgm_array
         return self._input_data
@@ -86,8 +73,8 @@ def create_grid_from_input_data(self, check_ids: bool = True) -> Grid:
 
         Returns a Grid object with the arrays filled with the PowerGridModel input.
         """
-        for pgm_name in PGM_ARRAYS:
-            if pgm_name in self._input_data:
+        for pgm_name in ComponentType:
+            if pgm_name in self._input_data and hasattr(self.grid, pgm_name):
                 pgm_ds_array_class = getattr(self.grid, pgm_name).__class__
                 pgm_ds_array = pgm_ds_array_class(self._input_data[pgm_name])
                 self.grid.append(pgm_ds_array, check_max_id=False)
@@ -155,12 +142,13 @@ def update_grid(self) -> None:
         """
         if not self.output_data:
             raise PGMCoreException("Can not update grid without output_data")
-        for array_name in PGM_ARRAYS:
-            if array_name in self.output_data.keys():
-                internal_array = getattr(self.grid, array_name)
-                pgm_output_array = self.output_data[array_name]
-                fields = self._match_dtypes(pgm_output_array.dtype, internal_array.dtype)
-                internal_array[fields] = pgm_output_array[fields]
+        for array_name in self.output_data.keys():
+            if not hasattr(self.grid, array_name):
+                continue
+            internal_array = getattr(self.grid, array_name)
+            pgm_output_array = self.output_data[array_name]
+            fields = self._match_dtypes(pgm_output_array.dtype, internal_array.dtype)
+            internal_array[fields] = pgm_output_array[fields]
 
     @staticmethod
     def _match_dtypes(first_dtype: np.dtype, second_dtype: np.dtype):

src/power_grid_model_ds/errors.py

@@ -2,7 +2,6 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 
-from power_grid_model_ds._core.load_flow import PGMCoreException
 from power_grid_model_ds._core.model.arrays.base.errors import (
     ArrayDefinitionError,
     MultipleRecordsReturned,
@@ -14,6 +13,7 @@
     MissingNodeError,
     NoPathBetweenNodes,
 )
+from power_grid_model_ds._core.power_grid_model_interface import PGMCoreException
 
 __all__ = [
     "PGMCoreException",

tests/integration/loadflow/test_power_grid_model.py

@@ -3,20 +3,23 @@
 # SPDX-License-Identifier: MPL-2.0
 
 
+from enum import StrEnum
+from unittest.mock import patch
+
 import numpy as np
 import pytest
-from power_grid_model import TapChangingStrategy, initialize_array
+from power_grid_model import ComponentType, TapChangingStrategy, initialize_array
 
 from power_grid_model_ds._core.data_source.generator.grid_generators import RadialGridGenerator
-from power_grid_model_ds._core.load_flow import PowerGridModelInterface
 from power_grid_model_ds._core.model.arrays import (
     LineArray,
     NodeArray,
     SourceArray,
     SymLoadArray,
 )
 from power_grid_model_ds._core.model.grids.base import Grid
-from tests.fixtures.arrays import DefaultedCustomLineArray, DefaultedCustomNodeArray
+from power_grid_model_ds._core.power_grid_model_interface import PowerGridModelInterface
+from tests.fixtures.arrays import DefaultedCustomNodeArray
 from tests.fixtures.grid_classes import ExtendedGrid, ExtendedGridNoDefaults
 
 # pylint: disable=missing-function-docstring,missing-class-docstring
@@ -111,27 +114,55 @@ def test_grid_with_automatic_tap_regulator(self, grid_with_tap_regulator: Grid):
         assert output["transformer_tap_regulator"]["tap_pos"][0] > 0
 
 
-class PowerGridModelInterfaceMethods:
-    def test_update_grid(self):
-        """Tests the power flow on a randomly configured grid and update grid with results"""
-        grid_generator = RadialGridGenerator(grid_class=Grid, nr_nodes=5, nr_sources=1, nr_nops=0)
-        grid = grid_generator.run(seed=0)
+@pytest.fixture
+def base_grid() -> Grid:
+    grid_generator = RadialGridGenerator(grid_class=Grid, nr_nodes=5, nr_sources=1, nr_nops=0)
+    return grid_generator.run(seed=0)
 
-        grid.node = DefaultedCustomNodeArray(grid.node.data)
-        grid.line = DefaultedCustomLineArray(grid.line.data)
 
-        core_interface = PowerGridModelInterface(grid=grid)
+class TestPowerGridModelInterfaceMethods:
+    @pytest.fixture
+    def extended_grid(self) -> ExtendedGrid:
+        grid_generator = RadialGridGenerator(grid_class=ExtendedGrid, nr_nodes=5, nr_sources=1, nr_nops=0)
+        return grid_generator.run(seed=0)
+
+    def test_update_grid(self, extended_grid):
+        """Tests the power flow on a randomly configured grid and update grid with results"""
+        core_interface = PowerGridModelInterface(grid=extended_grid)
+
         core_interface.create_input_from_grid()
         core_interface.calculate_power_flow()
         core_interface.update_grid()
 
         # voltage should be in neighbourhood of 10500
-        assert grid.node.u[0] == pytest.approx(10_500, 0.1)
-        assert grid.node.u[1] == pytest.approx(10_500, 0.1)
+        assert extended_grid.node.u[0] == pytest.approx(10_500, 0.1)
+        assert extended_grid.node.u[1] == pytest.approx(10_500, 0.1)
         # all lines have a current
-        assert all(grid.line.i_from > 0)
+        assert all(extended_grid.line.i_from > 0)
+
+    def test_create_input_from_grid_with_additional_component(self, extended_grid):
+        # If we delete the node array
+        del extended_grid.node
+        # The input_data should have a node key
+        core_interface = PowerGridModelInterface(grid=extended_grid)
+        input_data = core_interface.create_input_from_grid()
+        assert "node" not in input_data
+
+    def test_update_grid_with_additional_component(self, extended_grid):
+        core_interface = PowerGridModelInterface(grid=extended_grid)
 
-    def test_update_model(self):
+        core_interface.create_input_from_grid()
+        core_interface.calculate_power_flow()
+
+        # When we delete the node array from the grid
+        del core_interface.grid.node
+        core_interface.update_grid()
+
+        # We should not crash and still fill the line array with results
+        assert not hasattr(extended_grid, "node")
+        assert all(extended_grid.line.i_from > 0)
+
+    def test_update_model(self, extended_grid):
         """Test whether a pgm model can be updated and returns different results"""
         grid_generator = RadialGridGenerator(grid_class=Grid, nr_nodes=5, nr_sources=1, nr_nops=0)
         grid = grid_generator.run(seed=0)
@@ -182,7 +213,7 @@ def test_setup_model(self):
 
         core_interface = PowerGridModelInterface(grid=grid)
         assert core_interface.model is None
-        assert core_interface._input_data is None
+        assert core_interface._input_data == {}
         core_interface.setup_model()
         assert core_interface.model
         assert core_interface._input_data
@@ -324,3 +355,16 @@ def test_create_extended_grid_without_default_from_input_data(self, input_data_p
 
         with pytest.raises(ValueError, match="Missing required columns: {'u'}"):
             core_interface.create_grid_from_input_data()
+
+    def test_with_additional_component_type_in_input_data(self, input_data_pgm):
+        # If there is a new component from PGM in the input_data
+        input_data_pgm["new_component"] = initialize_array("input", "node", 2)
+        ExtendedComponentType = StrEnum(  # type: ignore
+            "ExtendedComponentType",
+            {component.name: component.value for component in ComponentType} | {"NEW_COMPONENT": "new_component"},
+        )
+
+        # We should skip it, but still succesfully generate a grid
+        with patch("power_grid_model_ds._core.power_grid_model_interface.ComponentType", ExtendedComponentType):
+            core_interface = PowerGridModelInterface(input_data=input_data_pgm)
+            core_interface.create_grid_from_input_data()

tests/unit/data_source/generator/test_grid_generators.py

@@ -12,10 +12,10 @@
 from power_grid_model_ds._core.data_source.generator.arrays.node import NodeGenerator
 from power_grid_model_ds._core.data_source.generator.arrays.source import SourceGenerator
 from power_grid_model_ds._core.data_source.generator.grid_generators import RadialGridGenerator
-from power_grid_model_ds._core.load_flow import PowerGridModelInterface
 from power_grid_model_ds._core.model.arrays import LineArray, NodeArray, SourceArray, SymLoadArray
 from power_grid_model_ds._core.model.graphs.models.base import BaseGraphModel
 from power_grid_model_ds._core.model.grids.base import Grid
+from power_grid_model_ds._core.power_grid_model_interface import PowerGridModelInterface
 
 
 def test_generate_random_grid():