feat: support all component types in interface

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"]

src/power_grid_model_ds/_core/load_flow.py → 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,7 +73,7 @@ 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:
+        for pgm_name in ComponentType:
             if pgm_name in self._input_data:
                 pgm_ds_array_class = getattr(self.grid, pgm_name).__class__
                 pgm_ds_array = pgm_ds_array_class(self._input_data[pgm_name])
@@ -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,19 +3,22 @@
 # 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 power_grid_model_ds._core.power_grid_model_interface import PowerGridModelInterface
 from tests.fixtures.arrays import ExtendedLineArray, ExtendedNodeArray
 from tests.fixtures.grid_classes import ExtendedGrid
 from tests.unit.model.grids.test_custom_grid import CustomGrid
@@ -112,20 +115,30 @@ def test_grid_with_automatic_tap_regulator(self, grid_with_tap_regulator: Grid):
         assert output["transformer_tap_regulator"]["tap_pos"][0] > 0
 
 
-class PowerGridModelInterfaceMethods:
+ExtendedComponentType = StrEnum(  # type: ignore
+    "ExtendedComponentType",
+    {component.name: component.value for component in ComponentType} | {"DUMMY_EXTRA_COMPONENT": "dummy"},
+)
+
+
+class TestPowerGridModelInterfaceMethods:
     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)
-
         grid.node = ExtendedNodeArray(grid.node.data)
         grid.line = ExtendedLineArray(grid.line.data)
 
-        core_interface = PowerGridModelInterface(grid=grid)
-        core_interface.create_input_from_grid()
-        core_interface.calculate_power_flow()
-        core_interface.update_grid()
+        # Patch ComponentType to include a value that is not present in the grid as an attribute
+        with patch("power_grid_model_ds._core.power_grid_model_interface.ComponentType", ExtendedComponentType):
+            core_interface = PowerGridModelInterface(grid=grid)
+            core_interface.create_input_from_grid()
+            core_interface.calculate_power_flow()
+            # Mock additional component type that is not in the grid
+            core_interface.output_data["dummy"] = np.array([])
+            core_interface.update_grid()
 
+        grid = core_interface.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)
@@ -183,7 +196,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

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():