feat: bypass pgm json conversion for simplicity

jaapschoutenalliander · jaapschoutenalliander · commit a6f0cbe74419 · 2025-10-22T15:11:07.000+02:00
Signed-off-by: jaapschoutenalliander &lt;jaap.schouten@alliander.com&gt;
diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py
@@ -7,155 +7,27 @@
 import dataclasses
 import json
 import logging
-from ast import literal_eval
 from pathlib import Path
 from typing import Dict, Optional
 
-import numpy as np
-from power_grid_model.utils import json_deserialize, json_serialize
-
-from power_grid_model_ds._core.load_flow import PGM_ARRAYS, PowerGridModelInterface
 from power_grid_model_ds._core.model.arrays.base.array import FancyArray
-
-# Constants
-EXTENDED_COLUMNS_KEY = "extended_columns"
-CUSTOM_ARRAYS_KEY = "custom_arrays"
-EXTENSIONS_KEY = "pgm_ds_extensions"
+from power_grid_model_ds._core.model.grids.base import Grid
 
 logger = logging.getLogger(__name__)
 
 
-def _extract_extensions_data(grid) -> Dict[str, Dict]:
-    """Extract extended columns and non-PGM arrays from a Grid object.
-
-    Args:
-        grid: The Grid object
-
-    Returns:
-        Dict containing extensions data with keys EXTENDED_COLUMNS_KEY and CUSTOM_ARRAYS_KEY
-    """
-    extensions: dict = {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}}
-
-    for field in dataclasses.fields(grid):
-        if field.name in ["graphs", "_id_counter"]:
-            continue
-
-        array = getattr(grid, field.name)
-        if not isinstance(array, FancyArray) or array.size == 0:
-            continue
-
-        array_name = field.name
-
-        if array_name in PGM_ARRAYS:
-            # Extract extended columns for PGM arrays
-            _extract_extended_columns(grid, array_name, array, extensions)
-        else:
-            # Store custom arrays not in PGM_ARRAYS
-            extensions[CUSTOM_ARRAYS_KEY][array_name] = {"dtype": str(array.dtype), "data": array.data.tolist()}
-
-    return extensions
-
-
-def _extract_extended_columns(grid, array_name: str, array: FancyArray, extensions: Dict) -> None:
-    """Extract extended columns from a PGM array."""
-    try:
-        interface = PowerGridModelInterface(grid=grid)
-        # pylint: disable=protected-access  # Accessing internal method for extension extraction
-        pgm_array = interface._create_power_grid_array(array_name)
-        pgm_columns = set(pgm_array.dtype.names or [])
-        ds_columns = set(array.columns)
-
-        # Find extended columns (columns in DS but not in PGM)
-        extended_cols = ds_columns - pgm_columns
-        if extended_cols:
-            extensions[EXTENDED_COLUMNS_KEY][array_name] = {col: array[col].tolist() for col in extended_cols}
-    except (AttributeError, KeyError, TypeError, ValueError) as e:
-        # Handle various failure modes:
-        # - KeyError: array_name not found in PGM arrays
-        # - AttributeError: array missing dtype/columns or interface method missing
-        # - TypeError/ValueError: invalid array configuration or data conversion issues
-        logger.warning(f"Failed to extract extensions for array '{array_name}': {e}")
-        extensions[CUSTOM_ARRAYS_KEY][array_name] = {"dtype": str(array.dtype), "data": array.data.tolist()}
-
-
-def _restore_extensions_data(grid, extensions_data: Dict) -> None:
-    """Restore extended columns and custom arrays to a Grid object.
-
-    Args:
-        grid: The Grid object to restore extensions to
-        extensions_data: Extensions data from _extract_extensions_data
-    """
-    # Restore extended columns
-    _restore_extended_columns(grid, extensions_data.get(EXTENDED_COLUMNS_KEY, {}))
-
-    # Restore custom arrays
-    _restore_custom_arrays(grid, extensions_data.get(CUSTOM_ARRAYS_KEY, {}))
-
-
-def _restore_extended_columns(grid, extended_columns: Dict) -> None:
-    """Restore extended columns to existing arrays."""
-    for array_name, extended_cols in extended_columns.items():
-        if not hasattr(grid, array_name):
-            logger.warning(f"Grid has no attribute '{array_name}' to restore")
-            continue
-
-        array = getattr(grid, array_name)
-        if not isinstance(array, FancyArray) or array.size == 0:
-            continue
-
-        for col_name, values in extended_cols.items():
-            # if hasattr(array, col_name):
-            try:
-                array[col_name] = values
-            except (AttributeError, IndexError, ValueError, TypeError) as e:
-                # Handle assignment failures:
-                # - IndexError: array size mismatch
-                # - ValueError/TypeError: incompatible data types
-                # - AttributeError: array doesn't support assignment
-                logger.warning(f"Failed to restore column '{col_name}' in array '{array_name}': {e}")
-
-
-def _parse_dtype(dtype_str: str) -> np.dtype:
-    """Parse a dtype string into a numpy dtype."""
-    if not isinstance(dtype_str, str):
-        raise ValueError(f"Invalid dtype string: {dtype_str}")
-
-    # Use numpy's dtype parsing - handle both eval-style and direct strings
-    if dtype_str.startswith("dtype("):
-        clean_dtype_str = dtype_str.replace("dtype(", "").replace(")", "")
-    else:
-        clean_dtype_str = dtype_str
-
-    # Use eval for complex dtype strings like "[('field', 'type'), ...]"
-    if clean_dtype_str.startswith("[") and clean_dtype_str.endswith("]"):
-        return np.dtype(literal_eval(clean_dtype_str))
-    return np.dtype(clean_dtype_str)
-
-
-def _construct_numpy_from_list(raw_data, dtype: np.dtype) -> np.ndarray:
-    """Construct a numpy array from a list with the specified dtype."""
-    if dtype.names:  # Structured dtype
-        # Convert from list of lists to list of tuples for structured array
-        if isinstance(raw_data[0], (list, tuple)) and len(raw_data[0]) == len(dtype.names):
-            data = np.array([tuple(row) for row in raw_data], dtype=dtype)
-        else:
-            data = np.array(raw_data, dtype=dtype)
-    else:
-        data = np.array(raw_data, dtype=dtype)
-    return data
-
-
-def _restore_custom_arrays(grid, custom_arrays: Dict) -> None:
+def _restore_grid_arrays(grid, custom_arrays: Dict) -> None:
     """Restore custom arrays to the grid."""
     for array_name, array_info in custom_arrays.items():
         if not hasattr(grid, array_name):
             continue
 
         try:
-            dtype = _parse_dtype(dtype_str=array_info["dtype"])
-            data = _construct_numpy_from_list(array_info["data"], dtype)
             array_field = grid.find_array_field(getattr(grid, array_name).__class__)
-            restored_array = array_field.type(data=data)
+            matched_columns = {
+                col: array_info["data"][col] for col in array_field.type().columns if col in array_info["data"]
+            }
+            restored_array = array_field.type(**matched_columns)
             setattr(grid, array_name, restored_array)
         except (AttributeError, KeyError, ValueError, TypeError) as e:
             # Handle restoration failures:
@@ -165,59 +37,35 @@ def _restore_custom_arrays(grid, custom_arrays: Dict) -> None:
             logger.warning(f"Failed to restore custom array '{array_name}': {e}")
 
 
-def _create_grid_from_input_data(input_data: Dict, target_grid_class=None):
-    """Create a Grid object from power-grid-model input data.
-
-    Args:
-        input_data: Power-grid-model input data
-        target_grid_class: Optional Grid class to create. If None, uses default Grid.
-
-    Returns:
-        Grid object populated with the input data
-    """
-    if target_grid_class is not None:
-        # Create empty grid of target type and populate it with input data
-        target_grid = target_grid_class.empty()
-        interface = PowerGridModelInterface(grid=target_grid, input_data=input_data)
-        return interface.create_grid_from_input_data()
-
-    # Use default Grid type
-    interface = PowerGridModelInterface(input_data=input_data)
-    return interface.create_grid_from_input_data()
-
-
 def save_grid_to_json(
     grid,
     path: Path,
-    use_compact_list: bool = True,
     indent: Optional[int] = None,
-    preserve_extensions: bool = True,
 ) -> 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
-        use_compact_list: Whether to use compact list format
         indent: JSON indentation (None for compact, positive int for indentation)
-        preserve_extensions: Whether to save extended columns and custom arrays
     Returns:
         Path: The path where the file was saved
     """
     path.parent.mkdir(parents=True, exist_ok=True)
 
-    # Convert Grid to power-grid-model input format and serialize
-    interface = PowerGridModelInterface(grid=grid)
-    input_data = interface.create_input_from_grid()
+    serialized_data = {}
+    for field in dataclasses.fields(grid):
+        if field.name in ["graphs", "_id_counter"]:
+            continue
 
-    core_data = json_serialize(input_data, use_compact_list=use_compact_list)
+        array = getattr(grid, field.name)
+        if not isinstance(array, FancyArray) or array.size == 0:
+            continue
 
-    # Parse and add extensions if requested
-    serialized_data = json.loads(core_data)
-    if preserve_extensions:
-        extensions = _extract_extensions_data(grid)
-        if extensions[EXTENDED_COLUMNS_KEY] or extensions[CUSTOM_ARRAYS_KEY]:
-            serialized_data[EXTENSIONS_KEY] = extensions
+        array_name = field.name
+        serialized_data[array_name] = {
+            "data": {name: array[name].tolist() for name in array.dtype.names},
+        }
 
     # Write to file
     with open(path, "w", encoding="utf-8") as f:
@@ -237,14 +85,13 @@ def load_grid_from_json(path: Path, target_grid_class=None):
         Grid: The deserialized Grid object of the specified target class
     """
     with open(path, "r", encoding="utf-8") as f:
-        data = json.load(f)
+        input_data = json.load(f)
 
-    # Extract extensions and deserialize core data
-    extensions = data.pop(EXTENSIONS_KEY, {EXTENDED_COLUMNS_KEY: {}, CUSTOM_ARRAYS_KEY: {}})
-    input_data = json_deserialize(json.dumps(data))
+    if target_grid_class is None:
+        target_grid = Grid.empty()
+    else:
+        target_grid = target_grid_class.empty()
 
-    # Create grid and restore extensions
-    grid = _create_grid_from_input_data(input_data, target_grid_class)
-    _restore_extensions_data(grid, extensions)
+    _restore_grid_arrays(target_grid, input_data)
 
-    return grid
+    return target_grid
diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py
@@ -16,8 +16,6 @@
 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 (
-    _extract_extensions_data,
-    _restore_extensions_data,
     load_grid_from_json,
     save_grid_to_json,
 )
@@ -48,6 +46,9 @@ class ExtendedGrid(Grid):
     node: ExtendedNodeArray
     line: ExtendedLineArray
 
+    # value_extension: float = 0.0
+    # dict_extension: dict = dict()
+
 
 @pytest.fixture
 def temp_dir():
@@ -91,11 +92,10 @@ def extended_grid():
 class TestSerializationFormats:
     """Test serialization across different formats and configurations"""
 
-    @pytest.mark.parametrize("preserve_ext", [(True), (False)])
-    def test_basic_serialization_roundtrip(self, basic_grid: Grid, temp_dir: Path, preserve_ext: bool):
+    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, preserve_extensions=preserve_ext)
+        result_path = save_grid_to_json(basic_grid, path)
         assert result_path.exists()
 
         # Load and verify
@@ -108,7 +108,7 @@ def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tem
         """Test extended serialization preserving custom data"""
         path = temp_dir / "extended.json"
 
-        save_grid_to_json(extended_grid, path, preserve_extensions=True)
+        save_grid_to_json(extended_grid, path)
         loaded_grid = load_grid_from_json(path, target_grid_class=ExtendedGrid)
 
         # Verify core data
@@ -140,7 +140,7 @@ def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir:
         path = temp_dir / "extended.json"
 
         # Save extended grid
-        save_grid_to_json(extended_grid, path, preserve_extensions=True)
+        save_grid_to_json(extended_grid, path)
         loaded_grid = load_grid_from_json(path, target_grid_class=Grid)
 
         # Core data should transfer
@@ -151,22 +151,6 @@ def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, temp_dir:
 class TestExtensionHandling:
     """Test extension data handling and edge cases"""
 
-    def test_missing_extension_keys(self):
-        """Test graceful handling of missing extension keys"""
-        basic_grid = Grid.empty()
-
-        # Test various malformed extension data
-        test_cases = [
-            {},  # Empty
-            {"extended_columns": {}},  # Missing custom_arrays
-            {"custom_arrays": {}},  # Missing extended_columns
-            {"extended_columns": {"test": "value"}},  # Invalid structure
-        ]
-
-        for extensions in test_cases:
-            # Should not raise
-            _restore_extensions_data(basic_grid, extensions)
-
     def test_custom_array_serialization_roundtrip(self, temp_dir: Path):
         """Test serialization and loading of grids with custom arrays"""
 
@@ -198,7 +182,7 @@ class GridWithCustomArray(Grid):
 
         # Test JSON serialization
         json_path = temp_dir / "custom_array.json"
-        save_grid_to_json(grid, json_path, preserve_extensions=True)
+        save_grid_to_json(grid, json_path)
 
         # Load back and verify
         loaded_grid = load_grid_from_json(json_path, target_grid_class=GridWithCustomArray)
@@ -231,31 +215,31 @@ def test_empty_grid_handling(self, temp_dir: Path):
         loaded_json = load_grid_from_json(json_path, target_grid_class=Grid)
         assert loaded_json.node.size == 0
 
-    def test_custom_array_extraction_edge_cases(self, temp_dir: Path):
-        """Test edge cases in custom array extraction"""
-        # Test with grid that has complex custom arrays that might cause extraction issues
-        extended_grid = ExtendedGrid.empty()
+    # def test_custom_array_extraction_edge_cases(self, temp_dir: Path):
+    #     """Test edge cases in custom array extraction"""
+    #     # Test with grid that has complex custom arrays that might cause extraction issues
+    #     extended_grid = ExtendedGrid.empty()
 
-        # Add data that might cause issues during extraction
-        nodes = ExtendedNodeArray(
-            id=[1, 2],
-            u_rated=[10000, 10000],
-            u=[float("nan"), float("inf")],  # Edge case values
-        )
-        extended_grid.append(nodes)
+    #     # Add data that might cause issues during extraction
+    #     nodes = ExtendedNodeArray(
+    #         id=[1, 2],
+    #         u_rated=[10000, 10000],
+    #         u=[float("nan"), float("inf")],  # Edge case values
+    #     )
+    #     extended_grid.append(nodes)
 
-        # Should handle edge case values gracefully
-        extensions = _extract_extensions_data(extended_grid)
-        assert "extended_columns" in extensions
-        assert "custom_arrays" in extensions
+    #     # Should handle edge case values gracefully
+    #     extensions = _extract_extensions_data(extended_grid)
+    #     assert "extended_columns" in extensions
+    #     assert "custom_arrays" in extensions
 
-        # Test saving and loading with these edge cases
-        json_path = temp_dir / "edge_cases.json"
-        save_grid_to_json(extended_grid, json_path, preserve_extensions=True)
+    #     # Test saving and loading with these edge cases
+    #     json_path = temp_dir / "edge_cases.json"
+    #     save_grid_to_json(extended_grid, json_path, preserve_extensions=True)
 
-        # Should load without issues
-        loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid)
-        assert loaded_grid.node.size == 2
+    #     # Should load without issues
+    #     loaded_grid = load_grid_from_json(json_path, target_grid_class=Grid)
+    #     assert loaded_grid.node.size == 2
 
     def test_invalid_extension_data_recovery(self, temp_dir: Path):
         """Test recovery from invalid extension data"""
@@ -265,7 +249,7 @@ def test_invalid_extension_data_recovery(self, temp_dir: Path):
         extended_grid.append(nodes)
 
         json_path = temp_dir / "test_recovery.json"
-        save_grid_to_json(extended_grid, json_path, preserve_extensions=True)
+        save_grid_to_json(extended_grid, json_path)
 
         # Corrupt extension data
         with open(json_path, "r") as f:
