diff --git a/src/power_grid_model_ds/__init__.py b/src/power_grid_model_ds/__init__.py index 78190b0..f1d7e6f 100644 --- a/src/power_grid_model_ds/__init__.py +++ b/src/power_grid_model_ds/__init__.py @@ -6,4 +6,8 @@ from power_grid_model_ds._core.model.graphs.container import GraphContainer from power_grid_model_ds._core.model.grids.base import Grid -__all__ = ["Grid", "GraphContainer", "PowerGridModelInterface"] +__all__ = [ + "Grid", + "GraphContainer", + "PowerGridModelInterface", +] 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 aacd636..6f4bf97 100644 --- a/src/power_grid_model_ds/_core/model/grids/base.py +++ b/src/power_grid_model_ds/_core/model/grids/base.py @@ -42,6 +42,7 @@ from power_grid_model_ds._core.model.grids._text_sources import TextSource from power_grid_model_ds._core.model.grids.helpers import set_feeder_ids, set_is_feeder from power_grid_model_ds._core.utils.pickle import get_pickle_path, load_from_pickle, save_to_pickle +from power_grid_model_ds._core.utils.serialization import load_grid_from_json, save_grid_to_json from power_grid_model_ds._core.utils.zip import file2gzip Self = TypeVar("Self", bound="Grid") @@ -360,7 +361,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 serialize() for better + interoperability and standardized format. Args: cache_dir (Path): The directory to save the cache to. @@ -435,6 +439,22 @@ def from_txt_file(cls, txt_file_path: Path): txt_lines = f.readlines() return TextSource(grid_class=cls).load_from_txt(*txt_lines) + def serialize(self, path: Path, **kwargs) -> Path: + """Serialize the grid to JSON format. + + Args: + path: Destination file path to write JSON to. + **kwargs: Additional keyword arguments forwarded to ``json.dump`` + Returns: + Path: The path where the file was saved. + """ + return save_grid_to_json(grid=self, path=path, **kwargs) + + @classmethod + def deserialize(cls: Type[Self], path: Path) -> Self: + """Deserialize the grid from JSON format.""" + return load_grid_from_json(path=path, target_grid_class=cls) + def set_feeder_ids(self): """Sets feeder and substation id properties in the grids arrays""" set_is_feeder(grid=self) diff --git a/src/power_grid_model_ds/_core/utils/serialization.py b/src/power_grid_model_ds/_core/utils/serialization.py new file mode 100644 index 0000000..cf67669 --- /dev/null +++ b/src/power_grid_model_ds/_core/utils/serialization.py @@ -0,0 +1,108 @@ +# SPDX-FileCopyrightText: Contributors to the Power Grid Model project +# +# 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 TYPE_CHECKING, Type, TypeVar + +from power_grid_model_ds._core.model.arrays.base.array import FancyArray + +if TYPE_CHECKING: + # Import only for type checking to avoid circular imports at runtime + from power_grid_model_ds._core.model.grids.base import Grid + + G = TypeVar("G", bound=Grid) +else: + # Runtime: don't import Grid to avoid circular import; keep unbound TypeVar + G = TypeVar("G") + +logger = logging.getLogger(__name__) + + +def save_grid_to_json(grid, path: Path, strict: bool = True, **kwargs) -> 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 + strict: Whether to raise an error if the grid object is not serializable. + **kwargs: Keyword arguments forwarded to json.dump (for example, indent, sort_keys, + ensure_ascii, etc.). + 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, FancyArray): + serialized_data[field.name] = { + "data": {name: field_value[name].tolist() for name in field_value.dtype.names}, + } + continue + + try: + json.dumps(field_value) + except TypeError as e: + if strict: + raise + logger.warning(f"Failed to serialize '{field.name}': {e}") + + # Write to file + with open(path, "w", encoding="utf-8") as f: + json.dump(serialized_data, f, **kwargs) + + return path + + +def load_grid_from_json(path: Path, target_grid_class: Type[G]) -> G: + """Load a Grid object from JSON format with cross-type loading support. + + Args: + path: The file path to load from + target_grid_class: Grid class to load into. + + 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) + + grid = target_grid_class.empty() + _restore_grid_values(grid, input_data) + graph_class = grid.graphs.__class__ + grid.graphs = graph_class.from_arrays(grid) + return grid + + +def _restore_grid_values(grid: G, json_data: dict) -> None: + """Restore arrays to the grid.""" + for attr_name, attr_values in json_data.items(): + if not hasattr(grid, attr_name): + logger.warning(f"Unexpected attribute '{attr_name}'") + continue + + grid_attr = getattr(grid, attr_name) + attr_class = grid_attr.__class__ + if isinstance(grid_attr, FancyArray): + if extra := set(attr_values["data"]) - set(grid_attr.columns): + logger.warning(f"{attr_name} has extra columns: {extra}") + + matched_columns = {col: attr_values["data"][col] for col in grid_attr.columns if col in attr_values["data"]} + restored_array = attr_class(**matched_columns) + setattr(grid, attr_name, restored_array) + continue + + # load other values + setattr(grid, attr_name, attr_class(attr_values)) diff --git a/tests/integration/visualizer_tests.py b/tests/integration/visualizer_tests.py index 2cb437a..9dbabdf 100644 --- a/tests/integration/visualizer_tests.py +++ b/tests/integration/visualizer_tests.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: MPL-2.0 from dataclasses import dataclass +from pathlib import Path from power_grid_model_ds import Grid from power_grid_model_ds._core.visualizer.app import visualize @@ -50,6 +51,11 @@ def visualize_grid_with_links(): if __name__ == "__main__": + r_grid = get_radial_grid() + r_grid.serialize(Path("json_path")) + + new_grid = Grid.deserialize(Path("json_path")) + visualize_grid() # visualize_coordinated_grid() # visualize_grid_with_links() diff --git a/tests/unit/utils/test_serialization.py b/tests/unit/utils/test_serialization.py new file mode 100644 index 0000000..c117ad8 --- /dev/null +++ b/tests/unit/utils/test_serialization.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Contributors to the Power Grid Model project +# +# 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 + +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 +from power_grid_model_ds.fancypy import array_equal + + +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" + + +@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 TestSerializationRoundtrips: + """Test serialization across different formats and configurations""" + + def test_basic_serialization_roundtrip(self, basic_grid: Grid, tmp_path: Path): + """Test basic serialization roundtrip for all formats""" + path = tmp_path / "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) + array_equal(loaded_grid.node, basic_grid.node) + array_equal(loaded_grid.line, basic_grid.line) + assert list(loaded_grid.node.id) == list(basic_grid.node.id) + + def test_extended_serialization_roundtrip(self, extended_grid: ExtendedGrid, tmp_path: Path): + """Test extended serialization preserving custom data""" + path = tmp_path / "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 + + # 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, tmp_path: Path): + """Test serialization of empty grids""" + empty_grid = Grid.empty() + + json_path = tmp_path / "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, tmp_path: Path): + """Test loading basic grid into extended type""" + path = tmp_path / "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 + array_equal(loaded_grid.node, basic_grid.node) + array_equal(loaded_grid.line, basic_grid.line) + + def test_extended_to_basic_loading(self, extended_grid: ExtendedGrid, tmp_path: Path): + """Test loading extended grid into basic type""" + path = tmp_path / "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 + array_equal(loaded_grid.node, extended_grid.node) + array_equal(loaded_grid.line, extended_grid.line) + + +class TestExtensionHandling: + """Test extension data handling and edge cases""" + + def test_custom_array_serialization_roundtrip(self, tmp_path: 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 = tmp_path / "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])