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,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,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,26 +115,48 @@ 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("ExtendedComponentType", {x.name: x.value for x in ComponentType} | {"onzin": "onzin"})
+
+
+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()
+        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()
+            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)
         # all lines have a current
         assert all(grid.line.i_from > 0)
 
+    # def test_update_grid_missing_component(self, core_interface_after_loadflow: PowerGridModelInterface):
+    #     """Test that update grid works when a component is missing in the grid."""
+    #     del core_interface_after_loadflow.grid.source
+
+    #     core_interface_after_loadflow.update_grid()
+
+    #     # voltage should be in neighbourhood of 10500
+    #     grid = core_interface_after_loadflow.grid
+    #     assert not all(grid.node.is_empty("u"))
+    #     assert not hasattr(grid, "source")
+
+    def test_input_from_grid_missing_component(self, grid: Grid):
+        # Mimic no source in grid
+        del grid.source
+
+        core_interface = PowerGridModelInterface(grid=grid)
+        assert "source" not in core_interface.create_input_from_grid()
+
     def test_update_model(self):
         """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)
@@ -183,7 +208,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():