Feat: add from_extended methods (#76)

* Add FancyArray.from_extended

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

* switch extended arrays

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

* fix tests

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

* Add Grid.from_extended

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

* cleanup

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

* rewrite to append method and add asserts to test

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

* Separate array_equal_with_nan

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

* Improve test

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

* Add test for array_equal_with_nan

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

* Apply suggestion from review

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

---------

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

Author: Thijss

src/power_grid_model_ds/_core/fancypy.py

@@ -8,6 +8,8 @@
 
 import numpy as np
 
+from power_grid_model_ds._core.utils.misc import array_equal_with_nan
+
 if TYPE_CHECKING:
     from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 
@@ -44,23 +46,5 @@ def sort(array: "FancyArray", axis=-1, kind=None, order=None) -> "FancyArray":
 def array_equal(array1: "FancyArray", array2: "FancyArray", equal_nan: bool = True) -> bool:
     """Return True if two arrays are equal."""
     if equal_nan:
-        return _array_equal_with_nan(array1, array2)
+        return array_equal_with_nan(array1.data, array2.data)
     return np.array_equal(array1.data, array2.data)
-
-
-def _array_equal_with_nan(array1: "FancyArray", array2: "FancyArray") -> bool:
-    # np.array_equal does not work with NaN values in structured arrays, so we need to compare column by column.
-    # related issue: https://github.com/numpy/numpy/issues/21539
-
-    if array1.columns != array2.columns:
-        return False
-
-    for column in array1.columns:
-        column_dtype = array1.dtype[column]
-        if np.issubdtype(column_dtype, np.str_):
-            if not np.array_equal(array1[column], array2[column]):
-                return False
-            continue
-        if not np.array_equal(array1[column], array2[column], equal_nan=True):
-            return False
-    return True

src/power_grid_model_ds/_core/model/arrays/base/array.py

@@ -323,3 +323,11 @@ def as_df(self: Self):
         if pandas is None:
             raise ImportError("pandas is not installed")
         return pandas.DataFrame(self._data)
+
+    @classmethod
+    def from_extended(cls: Type[Self], extended: Self) -> Self:
+        """Create an instance from an extended array."""
+        if not isinstance(extended, cls):
+            raise TypeError(f"Extended array must be of type {cls.__name__}, got {type(extended).__name__}")
+        dtype = cls.get_dtype()
+        return cls(data=np.array(extended[list(dtype.names)], dtype=dtype))

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

@@ -440,6 +440,24 @@ def set_feeder_ids(self):
         set_is_feeder(grid=self)
         set_feeder_ids(grid=self)
 
+    @classmethod
+    def from_extended(cls, extended: "Grid") -> "Grid":
+        """Create a grid from an extended Grid object."""
+        new_grid = cls.empty()
+
+        # Add nodes first, so that branches can reference them
+        new_grid.append(new_grid.node.__class__.from_extended(extended.node))
+
+        for field in dataclasses.fields(cls):
+            if field.name == "node":
+                continue  # already added
+            if issubclass(field.type, FancyArray):
+                extended_array = getattr(extended, field.name)
+                new_array = field.type.from_extended(extended_array)
+                new_grid.append(new_array, check_max_id=False)
+
+        return new_grid
+
 
 def _add_branch_array(branch: BranchArray | Branch3Array, grid: Grid):
     """Add a branch array to the grid"""

src/power_grid_model_ds/_core/utils/misc.py

@@ -39,3 +39,24 @@ def get_inherited_attrs(cls: Type, *private_attributes):
             retrieved_attributes[private_attr] = attr_dict
 
     return retrieved_attributes
+
+
+def array_equal_with_nan(array1: np.ndarray, array2: np.ndarray) -> bool:
+    """Compare two structured arrays for equality, treating NaN values as equal.
+
+    np.array_equal does not work with NaN values in structured arrays, so we need to compare column by column.
+    related issue: https://github.com/numpy/numpy/issues/21539
+    """
+    if array1.dtype.names != array2.dtype.names:
+        return False
+
+    columns: Sequence[str] = array1.dtype.names
+    for column in columns:
+        column_dtype = array1.dtype[column]
+        if np.issubdtype(column_dtype, np.str_):
+            if not np.array_equal(array1[column], array2[column]):
+                return False
+            continue
+        if not np.array_equal(array1[column], array2[column], equal_nan=True):
+            return False
+    return True

tests/fixtures/arrays.py

@@ -5,6 +5,7 @@
 import numpy as np
 from numpy._typing import NDArray
 
+from power_grid_model_ds._core.model.arrays import LineArray, NodeArray
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 from power_grid_model_ds._core.model.dtypes.sensors import NDArray3
 
@@ -57,3 +58,19 @@ class FancyTestArray3(FancyArray):
 
     test_float1: NDArray3[np.float64]
     test_float2: NDArray3[np.float64]
+
+
+class ExtendedNodeArray(NodeArray):
+    """Extends the node array with an output value"""
+
+    _defaults = {"u": 0}
+
+    u: NDArray[np.float64]
+
+
+class ExtendedLineArray(LineArray):
+    """Extends the line array with an output value"""
+
+    _defaults = {"i_from": 0}
+
+    i_from: NDArray[np.float64]

tests/fixtures/grid_classes.py

@@ -5,10 +5,13 @@
 from dataclasses import dataclass
 
 from power_grid_model_ds._core.model.grids.base import Grid
+from tests.fixtures.arrays import ExtendedLineArray, ExtendedNodeArray
 
 
 @dataclass
 class ExtendedGrid(Grid):
-    """Grid with an extra container"""
+    """ExtendedGrid class for testing purposes."""
 
+    node: ExtendedNodeArray
+    line: ExtendedLineArray
     extra_value: int = 123

tests/fixtures/grids.py

@@ -8,12 +8,10 @@
 
 from power_grid_model_ds._core.model.arrays import (
     LineArray,
-    LinkArray,
     NodeArray,
     SourceArray,
     SymLoadArray,
     ThreeWindingTransformerArray,
-    TransformerArray,
 )
 from power_grid_model_ds._core.model.enums.nodes import NodeType
 from power_grid_model_ds._core.model.grids.base import Grid
@@ -44,18 +42,18 @@ def build_basic_grid(grid: T) -> T:
     #                      ***
 
     # Add Substations
-    substation = NodeArray(id=[101], u_rated=[10_500.0], node_type=[NodeType.SUBSTATION_NODE.value])
+    substation = grid.node.__class__(id=[101], u_rated=[10_500.0], node_type=[NodeType.SUBSTATION_NODE.value])
     grid.append(substation, check_max_id=False)
 
     # Add Nodes
-    nodes = NodeArray(
+    nodes = grid.node.__class__(
         id=[102, 103, 104, 105, 106],
         u_rated=[10_500.0] * 4 + [400.0],
     )
     grid.append(nodes, check_max_id=False)
 
     # Add Lines
-    lines = LineArray(
+    lines = grid.line.__class__(
         id=[201, 202, 203, 204],
         from_status=[1, 1, 0, 1],
         to_status=[1, 1, 0, 1],
@@ -70,7 +68,7 @@ def build_basic_grid(grid: T) -> T:
     grid.append(lines, check_max_id=False)
 
     # Add a transformer
-    transformer = TransformerArray.empty(1)
+    transformer = grid.transformer.__class__.empty(1)
     transformer.id = 301
     transformer.from_status = 1
     transformer.to_status = 1
@@ -80,7 +78,7 @@ def build_basic_grid(grid: T) -> T:
     grid.append(transformer, check_max_id=False)
 
     # Add a link
-    link = LinkArray.empty(1)
+    link = grid.link.__class__.empty(1)
     link.id = 601
     link.from_status = 1
     link.to_status = 1
@@ -90,7 +88,7 @@ def build_basic_grid(grid: T) -> T:
     grid.append(link, check_max_id=False)
 
     # Loads
-    loads = SymLoadArray(
+    loads = grid.sym_load.__class__(
         id=[401, 402, 403, 404],
         node=[102, 103, 104, 105],
         type=[1] * 4,
@@ -101,7 +99,7 @@ def build_basic_grid(grid: T) -> T:
     grid.append(loads, check_max_id=False)
 
     # Add Source
-    source = SourceArray(id=[501], node=[101], status=[1], u_ref=[0.0])
+    source = grid.source.__class__(id=[501], node=[101], status=[1], u_ref=[0.0])
     grid.append(source, check_max_id=False)
     grid.check_ids()
 

tests/integration/loadflow/test_power_grid_model.py

@@ -23,27 +23,12 @@
 )
 from power_grid_model_ds._core.model.arrays.pgm_arrays import TransformerTapRegulatorArray
 from power_grid_model_ds._core.model.grids.base import Grid
+from tests.fixtures.arrays import ExtendedLineArray, ExtendedNodeArray
 from tests.unit.model.grids.test_custom_grid import CustomGrid
 
 # pylint: disable=missing-function-docstring,missing-class-docstring
 
 
-class ExtendedNodeArray(NodeArray):
-    """Extends the node array with an output value"""
-
-    _defaults = {"u": 0}
-
-    u: NDArray[np.float64]
-
-
-class ExtendedLineArray(LineArray):
-    """Extends the line array with an output value"""
-
-    _defaults = {"i_from": 0}
-
-    i_from: NDArray[np.float64]
-
-
 def test_load_flow_on_random():
     """Tests the power flow on a randomly configured grid"""
     grid_generator = RadialGridGenerator(grid_class=Grid, nr_nodes=5, nr_sources=1, nr_nops=0)

tests/unit/model/arrays/test_array.py

@@ -11,10 +11,11 @@
 
 from power_grid_model_ds._core import fancypy as fp
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
-from power_grid_model_ds._core.model.arrays.pgm_arrays import TransformerArray
+from power_grid_model_ds._core.model.arrays.pgm_arrays import LineArray, TransformerArray
 from power_grid_model_ds._core.model.constants import EMPTY_ID, empty
+from power_grid_model_ds._core.utils.misc import array_equal_with_nan
 from tests.conftest import FancyTestArray
-from tests.fixtures.arrays import FancyTestArray3
+from tests.fixtures.arrays import ExtendedLineArray, FancyTestArray3
 
 # pylint: disable=missing-function-docstring
 
@@ -289,3 +290,16 @@ def test_overflow_value():
     with pytest.raises(OverflowError):
         transformer.tap_min = -167
     assert transformer.tap_min == -128
+
+
+def test_from_extended_array():
+    extended_array = ExtendedLineArray.empty(3)
+    extended_array.id = [1, 2, 3]
+    extended_array.from_node = [4, 5, 6]
+    extended_array.to_node = [7, 8, 9]
+    extended_array.from_status = [1, 0, 1]
+    extended_array.from_status = [0, 1, 0]
+
+    array = LineArray.from_extended(extended_array)
+    assert not isinstance(array, ExtendedLineArray)
+    array_equal_with_nan(array.data, extended_array[array.columns])

tests/unit/model/grids/test_grid_base.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import pytest
+from numpy.ma.testutils import assert_array_equal
 
 from power_grid_model_ds._core.model.arrays import (
     LineArray,
@@ -20,6 +21,7 @@
 from power_grid_model_ds._core.model.constants import EMPTY_ID
 from power_grid_model_ds._core.model.grids.base import Grid
 from tests.fixtures.grid_classes import ExtendedGrid
+from tests.fixtures.grids import build_basic_grid
 
 # pylint: disable=missing-function-docstring,missing-class-docstring
 
@@ -50,6 +52,20 @@ def test_initialize_empty_extended_grid():
     assert isinstance(grid, ExtendedGrid)
 
 
+def test_from_extended_grid():
+    extended_grid = build_basic_grid(ExtendedGrid.empty())
+    grid = Grid.from_extended(extended_grid)
+    assert not isinstance(grid, ExtendedGrid)
+    assert_array_equal(grid.line.data, extended_grid.line.data[grid.line.columns])
+    assert grid.node.size
+    assert grid.branches.size
+    assert grid.graphs.active_graph.nr_nodes == len(grid.node)
+    assert grid.graphs.complete_graph.nr_nodes == len(grid.branches)
+
+    assert extended_grid.id_counter == grid.id_counter
+    assert extended_grid.max_id == grid.max_id
+
+
 def test_grid_build(basic_grid: Grid):
     grid = basic_grid