diff --git a/examples/development/entrypoint_plotting.py b/examples/development/entrypoint_plotting.py index 9503a3d78..b4b90fd9a 100644 --- a/examples/development/entrypoint_plotting.py +++ b/examples/development/entrypoint_plotting.py @@ -9,20 +9,28 @@ def main(): scenario_name = "singapore_geylang" # TODO: Split singapore_all scenario into separate folder - mesmo.utils.cleanup() results_path = mesmo.utils.get_results_path("run_operation_problem", scenario_name) + results_raw = mesmo.api.run_nominal_operation_problem( scenario_name, results_path=results_path, store_results=False, recreate_database=False ) results = results_raw.get_run_results() - # Roundtrip save/load to/from JSON, just for demonstration - with open(results_path / "results.json", "w", encoding="utf-8") as file: - print("Dumping results to file") - file.write(results.model_dump_json()) - with open(results_path / "results.json", "r", encoding="utf-8") as file: - print("Loading results from file") - results = mesmo.data_models.RunResults.model_validate_json(file.read()) + # # Roundtrip save/load to/from JSON, just for demonstration + # print("Writing results to file") + # results.to_json_file(results_path / "results.json") + # print("Loading results from file") + # results = mesmo.data_models.RunResults.from_json_file(results_path / "results.json") + + # Write results to compressed JSON + print("Writing results to file") + results.to_json_file(mesmo.config.base_path / "results" / "results.json.bz2", compress=True) + + # Load results from compressed JSON + print("Loading results from file") + results = mesmo.data_models.RunResults.from_json_file( + mesmo.config.base_path / "results" / "results.json.bz2", decompress=True + ) # Sample plotting to file, just for demonstration plots.plot_to_file(plots.der_active_power_time_series, results=results, results_path=results_path) diff --git a/mesmo/data_models/base_model.py b/mesmo/data_models/base_model.py index 947e9989d..dadafab2c 100644 --- a/mesmo/data_models/base_model.py +++ b/mesmo/data_models/base_model.py @@ -1,6 +1,8 @@ """Custom pydantic base model.""" -from typing import Annotated, Optional +import bz2 +import pathlib +from typing import Annotated, Optional, TypeVar import numpy as np import pandas as pd @@ -8,6 +10,8 @@ import pydantic as pyd from pydantic.json import timedelta_isoformat +Model = TypeVar("Model", bound="BaseModel") + class BaseModel(pyd.BaseModel): """Custom base model which overrides pydantic default definitions.""" @@ -23,6 +27,36 @@ class BaseModel(pyd.BaseModel): }, ) + def to_json_file(self, file_path: pathlib.Path, compress=False): + """Write data model to given file path as JSON file. + + Args: + file_path (pathlib.Path): File output path + """ + if compress: + with bz2.open(file_path, "wt", encoding="utf-8") as file: + file.write(self.model_dump_json()) + else: + with open(file_path, "w", encoding="utf-8") as file: + file.write(self.model_dump_json()) + + @classmethod + def from_json_file(cls: type[Model], file_path: pathlib.Path, decompress=False) -> Model: + """Load data model from given JSON file path. + + Args: + file_path (pathlib.Path): File input path + + Returns: + Instantiated data model based on given JSON file + """ + if decompress: + with bz2.open(file_path, "rt", encoding="utf-8") as file: + return cls.model_validate_json(file.read()) + else: + with open(file_path, "r", encoding="utf-8") as file: + return cls.model_validate_json(file.read()) + def get_index_annotation(dtype: type[pd.Timestamp | int | str], optional: bool = False) -> type[pd.Index]: """Get pydantic-compatible type annotation for pandas.Index."""