feat: grid serialization

src/power_grid_model_ds/__init__.py

@@ -5,5 +5,15 @@
 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.utils.serialization import (
+    load_grid_from_json,
+    save_grid_to_json,
+)
 
-__all__ = ["Grid", "GraphContainer", "PowerGridModelInterface"]
+__all__ = [
+    "Grid",
+    "GraphContainer",
+    "PowerGridModelInterface",
+    "save_grid_to_json",
+    "load_grid_from_json",
+]

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

@@ -360,7 +360,10 @@ def get_downstream_nodes(self, node_id: int, inclusive: bool = False):
         )
 
     def cache(self, cache_dir: Path, cache_name: str, compress: bool = True):
-        """Cache Grid to a folder
+        """Cache Grid to a folder using pickle format.
+
+        Note: Consider using save_to_json() for better
+        interoperability and standardized format.
 
         Args:
             cache_dir (Path): The directory to save the cache to.

src/power_grid_model_ds/_core/utils/serialization.py

@@ -0,0 +1,104 @@
+# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <[email protected]>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""Serialization utilities for Grid objects using power-grid-model serialization with extensions support."""
+
+import dataclasses
+import json
+import logging
+from pathlib import Path
+from typing import Dict, Optional
+
+from power_grid_model_ds._core.model.arrays.base.array import FancyArray
+from power_grid_model_ds._core.model.grids.base import Grid
+
+logger = logging.getLogger(__name__)
+
+
+def _restore_grid_values(grid, input_data: Dict) -> None:
+    """Restore arrays to the grid."""
+    for attr_name, attr_values in input_data.items():
+        if not hasattr(grid, attr_name):
+            continue
+
+        if not issubclass(getattr(grid, attr_name).__class__, FancyArray):
+            setattr(grid, attr_name, attr_values)
+            continue
+
+        try:
+            array_field = grid.find_array_field(getattr(grid, attr_name).__class__)
+            matched_columns = {
+                col: attr_values["data"][col] for col in array_field.type().columns if col in attr_values["data"]
+            }
+            restored_array = array_field.type(**matched_columns)
+            setattr(grid, attr_name, restored_array)
+        except (AttributeError, KeyError, ValueError, TypeError) as e:
+            # Handle restoration failures:
+            # - KeyError: missing "dtype" or "data" keys
+            # - ValueError/TypeError: invalid dtype string or data conversion
+            # - AttributeError: grid methods/attributes missing
+            logger.warning(f"Failed to restore '{attr_name}': {e}")
+
+
+def save_grid_to_json(
+    grid,
+    path: Path,
+    indent: Optional[int] = None,
+) -> Path:
+    """Save a Grid object to JSON format using power-grid-model serialization with extensions support.
+
+    Args:
+        grid: The Grid object to serialize
+        path: The file path to save to
+        indent: JSON indentation (None for compact, positive int for indentation)
+    Returns:
+        Path: The path where the file was saved
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    serialized_data = {}
+    for field in dataclasses.fields(grid):
+        if field.name in ["graphs", "_id_counter"]:
+            continue
+
+        field_value = getattr(grid, field.name)
+        if isinstance(field_value, (int, float, str, bool)):
+            serialized_data[field.name] = field_value
+
+        if not isinstance(field_value, FancyArray) or field_value.size == 0:
+            continue
+
+        array_name = field.name
+        serialized_data[array_name] = {
+            "data": {name: field_value[name].tolist() for name in field_value.dtype.names},
+        }
+
+    # Write to file
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(serialized_data, f, indent=indent if indent and indent > 0 else None)
+
+    return path
+
+
+def load_grid_from_json(path: Path, target_grid_class=None):
+    """Load a Grid object from JSON format with cross-type loading support.
+
+    Args:
+        path: The file path to load from
+        target_grid_class: Optional Grid class to load into. If None, uses default Grid.
+
+    Returns:
+        Grid: The deserialized Grid object of the specified target class
+    """
+    with open(path, "r", encoding="utf-8") as f:
+        input_data = json.load(f)
+
+    if target_grid_class is None:
+        target_grid = Grid.empty()
+    else:
+        target_grid = target_grid_class.empty()
+
+    _restore_grid_values(target_grid, input_data)
+
+    return target_grid

tests/unit/utils/test_serialization.py

@@ -0,0 +1,215 @@
+# SPDX-FileCopyrightText: Contributors to the Power Grid Model project <[email protected]>
+#
+# SPDX-License-Identifier: MPL-2.0
+
+"""Comprehensive unit tests for Grid serialization with power-grid-model compatibility."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+import pytest
+from numpy.typing import NDArray
+
+from power_grid_model_ds import Grid
+from power_grid_model_ds._core.model.arrays.base.array import FancyArray
+from power_grid_model_ds._core.utils.serialization import (
+    load_grid_from_json,
+    save_grid_to_json,
+)
+from power_grid_model_ds.arrays import LineArray
+from power_grid_model_ds.arrays import NodeArray as BaseNodeArray
+
+
+class ExtendedNodeArray(BaseNodeArray):
+    """Test array with extended columns"""
+
+    _defaults = {"u": 0.0, "analysis_flag": 0}
+    u: NDArray[np.float64]
+    analysis_flag: NDArray[np.int32]
+
+
+class ExtendedLineArray(LineArray):
+    """Test array with extended columns"""
+
+    _defaults = {"i_from": 0.0, "loading_factor": 0.0}
+    i_from: NDArray[np.float64]
+    loading_factor: NDArray[np.float64]
+
+
+@dataclass
+class ExtendedGrid(Grid):
+    """Test grid with extended arrays"""
+
+    node: ExtendedNodeArray
+    line: ExtendedLineArray
+
+    value_extension: float = 0.0
+    str_extension: str = "default"
+    complex_extension: list = None
+
+
+@pytest.fixture
+def temp_dir():
+    """Temporary directory fixture"""
+    with TemporaryDirectory() as tmp_dir:
+        yield Path(tmp_dir)
+
+
+@pytest.fixture
+def basic_grid():
+    """Basic grid fixture"""
+    return Grid.from_txt("1 2", "2 3", "S10 1")
+
+
+@pytest.fixture
+def extended_grid():
+    """Extended grid fixture with additional columns"""
+    grid = ExtendedGrid.empty()
+    nodes = ExtendedNodeArray(
+        id=[1, 2, 3], u_rated=[10500, 10500, 10500], u=[10450, 10400, 10350], analysis_flag=[1, 0, 1]
+    )
+    lines = ExtendedLineArray(
+        id=[10, 11],
+        from_node=[1, 2],
+        to_node=[2, 3],
+        from_status=[1, 1],
+        to_status=[1, 1],
+        r1=[0.1, 0.15],
+        x1=[0.2, 0.25],
+        c1=[1e-6, 1.2e-6],
+        tan1=[0.0, 0.0],
+        i_n=[400, 350],
+        i_from=[150.5, 120.3],
+        loading_factor=[0.75, 0.68],
+    )
+    grid.append(nodes)
+    grid.append(lines)
+    return grid
+
+
+class TestSerializationFormats:
+    """Test serialization across different formats and configurations"""
+
+    def test_basic_serialization_roundtrip(self, basic_grid: Grid, temp_dir: Path):
+        """Test basic serialization roundtrip for all formats"""
+        path = temp_dir / "test.json"
+        result_path = save_grid_to_json(basic_grid, path)
+        assert result_path.exists()
+
+        # Load and verify
+        loaded_grid = load_grid_from_json(path, target_grid_class=Grid)
+        assert loaded_grid.node.size == basic_grid.node.size
+        assert loaded_grid.line.size == basic_grid.line.size
+        assert list(loaded_grid.node.id) == list(basic_grid.node.id)
+
+    def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, temp_dir: Path):
+        """Test extended serialization preserving custom data"""
+        path = temp_dir / "extended.json"
+
+        save_grid_to_json(extended_grid, path)
+        loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid)
+
+        # Verify core data
+        assert loaded_grid.node.size == extended_grid.node.size
+        assert loaded_grid.line.size == extended_grid.line.size
+        assert loaded_grid.value_extension == extended_grid.value_extension
+        assert loaded_grid.str_extension == extended_grid.str_extension
+        assert loaded_grid.complex_extension is None
+
+        # Verify extended data
+        np.testing.assert_array_equal(loaded_grid.node.u, extended_grid.node.u)
+        np.testing.assert_array_equal(loaded_grid.line.i_from, extended_grid.line.i_from)
+
+    def test_empty_grid_handling(self, temp_dir: Path):
+        """Test serialization of empty grids"""
+        empty_grid = Grid.empty()
+
+        json_path = temp_dir / "empty.json"
+
+        # Should handle empty grids
+        save_grid_to_json(empty_grid, json_path)
+
+        # Should load back as empty
+        loaded_json = load_grid_from_json(json_path, target_grid_class=Grid)
+        assert loaded_json.node.size == 0
+
+
+class TestCrossTypeCompatibility:
+    """Test cross-type loading and compatibility"""
+
+    def test_basic_to_extended_loading(self, basic_grid: Grid, temp_dir: Path):
+        """Test loading basic grid into extended type"""
+        path = temp_dir / "basic.json"
+
+        # Save basic grid
+        save_grid_to_json(basic_grid, path)
+        loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid)
+
+        # Core data should transfer
+        assert loaded_grid.node.size == basic_grid.node.size
+        assert loaded_grid.line.size == basic_grid.line.size
+
+    def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir: Path):
+        """Test loading extended grid into basic type"""
+        path = temp_dir / "extended.json"
+
+        # Save extended grid
+        save_grid_to_json(extended_grid, path)
+        loaded_grid = load_grid_from_json(path, target_grid_class=Grid)
+
+        # Core data should transfer
+        assert loaded_grid.node.size == extended_grid.node.size
+        assert loaded_grid.line.size == extended_grid.line.size
+
+
+class TestExtensionHandling:
+    """Test extension data handling and edge cases"""
+
+    def test_custom_array_serialization_roundtrip(self, temp_dir: Path):
+        """Test serialization and loading of grids with custom arrays"""
+
+        # Create a custom array type that properly extends FancyArray
+        class CustomMetadataArray(FancyArray):
+            """Custom metadata array for testing"""
+
+            _defaults = {"metadata_value": 0.0, "category": 0}
+
+            id: NDArray[np.int32]
+            metadata_value: NDArray[np.float64]
+            category: NDArray[np.int32]
+
+        # Create a grid with custom arrays
+        @dataclass
+        class GridWithCustomArray(Grid):
+            custom_metadata: CustomMetadataArray
+
+        # Create test grid with custom data
+        grid = GridWithCustomArray.empty()
+
+        # Add some basic grid data
+        nodes = grid.node.__class__(id=[1, 2], u_rated=[10000, 10000])
+        grid.append(nodes)
+
+        # Add custom metadata
+        custom_data = CustomMetadataArray(id=[100, 200, 300], metadata_value=[1.5, 2.5, 3.5], category=[1, 2, 1])
+        grid.custom_metadata = custom_data
+
+        # Test JSON serialization
+        json_path = temp_dir / "custom_array.json"
+        save_grid_to_json(grid, json_path)
+
+        # Load back and verify
+        loaded_grid = load_grid_from_json(json_path, target_grid_class=GridWithCustomArray)
+
+        # Verify core data
+        assert loaded_grid.node.size == 2
+        np.testing.assert_array_equal(loaded_grid.node.id, [1, 2])
+
+        # Verify custom array was preserved
+        assert hasattr(loaded_grid, "custom_metadata")
+        assert loaded_grid.custom_metadata.size == 3
+        np.testing.assert_array_equal(loaded_grid.custom_metadata.id, [100, 200, 300])
+        np.testing.assert_array_almost_equal(loaded_grid.custom_metadata.metadata_value, [1.5, 2.5, 3.5])
+        np.testing.assert_array_equal(loaded_grid.custom_metadata.category, [1, 2, 1])