From 6bebd284865c4bb7e084435ff171faf7f0eaa8e7 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Wed, 16 Jul 2025 17:30:16 +0200
Subject: [PATCH 01/10] Add FancyArray.from_extended

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 .../_core/model/arrays/base/array.py                   | 10 ++++++++++
 tests/fixtures/arrays.py                               |  9 +++++++++
 tests/unit/model/arrays/test_array.py                  |  9 ++++++++-
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/power_grid_model_ds/_core/model/arrays/base/array.py b/src/power_grid_model_ds/_core/model/arrays/base/array.py
index 3ba132a..929f837 100644
--- a/src/power_grid_model_ds/_core/model/arrays/base/array.py
+++ b/src/power_grid_model_ds/_core/model/arrays/base/array.py
@@ -323,3 +323,13 @@ 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__}")
+        new_array = cls.zeros(len(extended), empty_id=False)
+        for column in cls.get_dtype().names:
+            new_array[column] = extended[column]
+        return new_array
diff --git a/tests/fixtures/arrays.py b/tests/fixtures/arrays.py
index ac23d25..0af0478 100644
--- a/tests/fixtures/arrays.py
+++ b/tests/fixtures/arrays.py
@@ -36,6 +36,15 @@ class DifferentFancyTestArray(FancyArray):
     test_bool2: NDArray[np.bool_]
 
 
+class ExtendedFancyTestArray(FancyTestArray):
+    """Test array with some extended attributes"""
+
+    test_int2: NDArray[np.int64]
+    test_float2: NDArray[np.float64]
+    test_str2: NDArray[np.str_]
+    test_bool2: NDArray[np.bool_]
+
+
 class LongColumnNameFancyTestArray(FancyArray):
     """Test array with a very long column name"""
 
diff --git a/tests/unit/model/arrays/test_array.py b/tests/unit/model/arrays/test_array.py
index 8868285..0d423b6 100644
--- a/tests/unit/model/arrays/test_array.py
+++ b/tests/unit/model/arrays/test_array.py
@@ -14,7 +14,7 @@
 from power_grid_model_ds._core.model.arrays.pgm_arrays import TransformerArray
 from power_grid_model_ds._core.model.constants import EMPTY_ID, empty
 from tests.conftest import FancyTestArray
-from tests.fixtures.arrays import FancyTestArray3
+from tests.fixtures.arrays import ExtendedFancyTestArray, FancyTestArray3
 
 # pylint: disable=missing-function-docstring
 
@@ -289,3 +289,10 @@ def test_overflow_value():
     with pytest.raises(OverflowError):
         transformer.tap_min = -167
     assert transformer.tap_min == -128
+
+
+def test_from_extended_array():
+    extended_array = ExtendedFancyTestArray.zeros(2)
+    array = FancyTestArray.from_extended(extended_array)
+    assert not isinstance(array, ExtendedFancyTestArray)
+    assert_array_equal(array.data, extended_array[array.columns])

From 72e3c614b035edfe9015a4eefda9980c5a5412ad Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Wed, 16 Jul 2025 17:34:37 +0200
Subject: [PATCH 02/10] switch extended arrays

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 tests/fixtures/arrays.py                      | 26 ++++++++++++-------
 .../loadflow/test_power_grid_model.py         | 17 +-----------
 tests/unit/model/arrays/test_array.py         | 10 +++----
 3 files changed, 23 insertions(+), 30 deletions(-)

diff --git a/tests/fixtures/arrays.py b/tests/fixtures/arrays.py
index 0af0478..63c176c 100644
--- a/tests/fixtures/arrays.py
+++ b/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
 
@@ -36,15 +37,6 @@ class DifferentFancyTestArray(FancyArray):
     test_bool2: NDArray[np.bool_]
 
 
-class ExtendedFancyTestArray(FancyTestArray):
-    """Test array with some extended attributes"""
-
-    test_int2: NDArray[np.int64]
-    test_float2: NDArray[np.float64]
-    test_str2: NDArray[np.str_]
-    test_bool2: NDArray[np.bool_]
-
-
 class LongColumnNameFancyTestArray(FancyArray):
     """Test array with a very long column name"""
 
@@ -66,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]
diff --git a/tests/integration/loadflow/test_power_grid_model.py b/tests/integration/loadflow/test_power_grid_model.py
index a91816e..a16b684 100644
--- a/tests/integration/loadflow/test_power_grid_model.py
+++ b/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)
diff --git a/tests/unit/model/arrays/test_array.py b/tests/unit/model/arrays/test_array.py
index 0d423b6..f1c6879 100644
--- a/tests/unit/model/arrays/test_array.py
+++ b/tests/unit/model/arrays/test_array.py
@@ -11,10 +11,10 @@
 
 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 tests.conftest import FancyTestArray
-from tests.fixtures.arrays import ExtendedFancyTestArray, FancyTestArray3
+from tests.fixtures.arrays import ExtendedLineArray, FancyTestArray3
 
 # pylint: disable=missing-function-docstring
 
@@ -292,7 +292,7 @@ def test_overflow_value():
 
 
 def test_from_extended_array():
-    extended_array = ExtendedFancyTestArray.zeros(2)
-    array = FancyTestArray.from_extended(extended_array)
-    assert not isinstance(array, ExtendedFancyTestArray)
+    extended_array = ExtendedLineArray.zeros(2)
+    array = LineArray.from_extended(extended_array)
+    assert not isinstance(array, ExtendedLineArray)
     assert_array_equal(array.data, extended_array[array.columns])

From 4e779fb5522d4806e460f13fea21182524770414 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Wed, 16 Jul 2025 17:44:23 +0200
Subject: [PATCH 03/10] fix tests

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 tests/fixtures/grid_classes.py |  5 ++++-
 tests/fixtures/grids.py        | 16 +++++++---------
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/tests/fixtures/grid_classes.py b/tests/fixtures/grid_classes.py
index c717003..4c4cd92 100644
--- a/tests/fixtures/grid_classes.py
+++ b/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."""
 
     extra_value: int = 123
+    node: ExtendedNodeArray
+    line: ExtendedLineArray
diff --git a/tests/fixtures/grids.py b/tests/fixtures/grids.py
index 842cb33..245488e 100644
--- a/tests/fixtures/grids.py
+++ b/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()
 

From b47c97e8c3b59e8e5c19f304d9d19f38b286ccc1 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Wed, 16 Jul 2025 18:01:22 +0200
Subject: [PATCH 04/10] Add Grid.from_extended

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 .../_core/model/grids/base.py                   | 17 +++++++++++++++++
 tests/unit/model/grids/test_grid_base.py        |  9 +++++++++
 2 files changed, 26 insertions(+)

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 bb8f35d..859cf79 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -440,6 +440,23 @@ def set_feeder_ids(self):
         set_is_feeder(grid=self)
         set_feeder_ids(grid=self)
 
+    @classmethod
+    def from_extended(cls, extended: "Grid", load_graphs: bool = True) -> "Grid":
+        """Create a grid from an extended Grid object."""
+        new_grid = cls.empty()
+
+        for field in dataclasses.fields(cls):
+            if issubclass(field.type, FancyArray):
+                extended_array = getattr(extended, field.name)
+                new_array = field.type.from_extended(extended_array)
+                setattr(new_grid, field.name, new_array)
+
+        new_grid._id_counter = new_grid.max_id
+
+        if load_graphs:
+            new_grid.graphs = GraphContainer.from_arrays(new_grid)
+        return new_grid
+
 
 def _add_branch_array(branch: BranchArray | Branch3Array, grid: Grid):
     """Add a branch array to the grid"""
diff --git a/tests/unit/model/grids/test_grid_base.py b/tests/unit/model/grids/test_grid_base.py
index ce3dfc8..e829f5e 100644
--- a/tests/unit/model/grids/test_grid_base.py
+++ b/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,13 @@ 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])
+
+
 def test_grid_build(basic_grid: Grid):
     grid = basic_grid
 

From 13fcb938e47feb07d47fd11249ef8f26141c3c48 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Wed, 16 Jul 2025 18:04:48 +0200
Subject: [PATCH 05/10] cleanup

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 tests/fixtures/grid_classes.py           | 2 +-
 tests/unit/model/grids/test_grid_base.py | 4 ++++
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/fixtures/grid_classes.py b/tests/fixtures/grid_classes.py
index 4c4cd92..05ffcca 100644
--- a/tests/fixtures/grid_classes.py
+++ b/tests/fixtures/grid_classes.py
@@ -12,6 +12,6 @@
 class ExtendedGrid(Grid):
     """ExtendedGrid class for testing purposes."""
 
-    extra_value: int = 123
     node: ExtendedNodeArray
     line: ExtendedLineArray
+    extra_value: int = 123
diff --git a/tests/unit/model/grids/test_grid_base.py b/tests/unit/model/grids/test_grid_base.py
index e829f5e..d66f436 100644
--- a/tests/unit/model/grids/test_grid_base.py
+++ b/tests/unit/model/grids/test_grid_base.py
@@ -57,6 +57,10 @@ def test_from_extended_grid():
     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)
 
 
 def test_grid_build(basic_grid: Grid):

From 0b90af25e57863887af73805b3de12b28c1a40f0 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Thu, 17 Jul 2025 10:03:06 +0200
Subject: [PATCH 06/10] rewrite to append method and add asserts to test

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 src/power_grid_model_ds/_core/model/grids/base.py | 13 +++++++------
 tests/unit/model/grids/test_grid_base.py          |  3 +++
 2 files changed, 10 insertions(+), 6 deletions(-)

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 859cf79..2dc64ed 100644
--- a/src/power_grid_model_ds/_core/model/grids/base.py
+++ b/src/power_grid_model_ds/_core/model/grids/base.py
@@ -441,20 +441,21 @@ def set_feeder_ids(self):
         set_feeder_ids(grid=self)
 
     @classmethod
-    def from_extended(cls, extended: "Grid", load_graphs: bool = True) -> "Grid":
+    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)
-                setattr(new_grid, field.name, new_array)
-
-        new_grid._id_counter = new_grid.max_id
+                new_grid.append(new_array, check_max_id=False)
 
-        if load_graphs:
-            new_grid.graphs = GraphContainer.from_arrays(new_grid)
         return new_grid
 
 
diff --git a/tests/unit/model/grids/test_grid_base.py b/tests/unit/model/grids/test_grid_base.py
index d66f436..3f8dc3f 100644
--- a/tests/unit/model/grids/test_grid_base.py
+++ b/tests/unit/model/grids/test_grid_base.py
@@ -62,6 +62,9 @@ def test_from_extended_grid():
     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

From c4714a3e8e5bceb2bb4eecf59acefe4abafb6c9a Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Thu, 17 Jul 2025 10:07:59 +0200
Subject: [PATCH 07/10] Separate array_equal_with_nan

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 src/power_grid_model_ds/_core/fancypy.py    | 22 +++------------------
 src/power_grid_model_ds/_core/utils/misc.py | 21 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 19 deletions(-)

diff --git a/src/power_grid_model_ds/_core/fancypy.py b/src/power_grid_model_ds/_core/fancypy.py
index a037e4f..ef4dc05 100644
--- a/src/power_grid_model_ds/_core/fancypy.py
+++ b/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
diff --git a/src/power_grid_model_ds/_core/utils/misc.py b/src/power_grid_model_ds/_core/utils/misc.py
index 0eeb64f..bf64df4 100644
--- a/src/power_grid_model_ds/_core/utils/misc.py
+++ b/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

From 0a40a2a6248c9a0a50387ffe0f1f1281715496b1 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Thu, 17 Jul 2025 10:13:23 +0200
Subject: [PATCH 08/10] Improve test

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 tests/unit/model/arrays/test_array.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests/unit/model/arrays/test_array.py b/tests/unit/model/arrays/test_array.py
index f1c6879..8e4b7c5 100644
--- a/tests/unit/model/arrays/test_array.py
+++ b/tests/unit/model/arrays/test_array.py
@@ -13,6 +13,7 @@
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
 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 ExtendedLineArray, FancyTestArray3
 
@@ -292,7 +293,13 @@ def test_overflow_value():
 
 
 def test_from_extended_array():
-    extended_array = ExtendedLineArray.zeros(2)
+    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)
-    assert_array_equal(array.data, extended_array[array.columns])
+    array_equal_with_nan(array.data, extended_array[array.columns])

From 3523024493d6a0b7f20a42f2e477fa6aa6d780fc Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Thu, 17 Jul 2025 10:32:05 +0200
Subject: [PATCH 09/10] Add test for array_equal_with_nan

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 tests/unit/utils/test_misc.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/tests/unit/utils/test_misc.py b/tests/unit/utils/test_misc.py
index 8081e2b..1864b54 100644
--- a/tests/unit/utils/test_misc.py
+++ b/tests/unit/utils/test_misc.py
@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from power_grid_model_ds._core.utils.misc import is_sequence
+from power_grid_model_ds._core.utils.misc import array_equal_with_nan, is_sequence
 
 # pylint: disable=missing-function-docstring
 
@@ -31,3 +31,9 @@ def test_dict_is_not_a_sequence():
 
 def test_string_is_not_a_sequence():
     assert not is_sequence("abc")
+
+
+def test_array_equal_with_nan():
+    array1 = np.array([(1, 2.0, "a"), (3, np.nan, "b")], dtype=[("col1", "i4"), ("col2", "f4"), ("col3", "U1")])
+    array2 = np.array([(1, 2.0, "a"), (3, np.nan, "b")], dtype=[("col1", "i4"), ("col2", "f4"), ("col3", "U1")])
+    assert array_equal_with_nan(array1, array2)

From 9ba4105889ba03d4f70268a52593df269741d153 Mon Sep 17 00:00:00 2001
From: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
Date: Thu, 17 Jul 2025 10:43:33 +0200
Subject: [PATCH 10/10] Apply suggestion from review

Signed-off-by: Thijs Baaijen <13253091+Thijss@users.noreply.github.com>
---
 src/power_grid_model_ds/_core/model/arrays/base/array.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/power_grid_model_ds/_core/model/arrays/base/array.py b/src/power_grid_model_ds/_core/model/arrays/base/array.py
index 929f837..e313dfd 100644
--- a/src/power_grid_model_ds/_core/model/arrays/base/array.py
+++ b/src/power_grid_model_ds/_core/model/arrays/base/array.py
@@ -329,7 +329,5 @@ 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__}")
-        new_array = cls.zeros(len(extended), empty_id=False)
-        for column in cls.get_dtype().names:
-            new_array[column] = extended[column]
-        return new_array
+        dtype = cls.get_dtype()
+        return cls(data=np.array(extended[list(dtype.names)], dtype=dtype))