Source code for r2x_core.system

"""R2X Core System class - subclass of infrasys.System with R2X-specific functionality."""

from collections.abc import Callable
from importlib.metadata import version
from pathlib import Path
from typing import Any

import orjson
from infrasys.component import Component
from infrasys.system import System as InfrasysSystem
from infrasys.utils.sqlite import backup
from loguru import logger

from . import units
from .utils import filter_kwargs_by_signatures
from .utils.files import get_r2x_cache_path


[docs] class System(InfrasysSystem): """R2X Core System class extending infrasys.System. Extends infrasys.System to provide R2X-specific functionality for data model translation and system construction. Adds convenience methods for component export and system manipulation. Parameters ---------- system_base : float | None, optional System base power in MVA for per-unit calculations. Default is None. name : str | None, optional Unique identifier for the system. Default is None. **kwargs Additional keyword arguments passed to infrasys.System (e.g., description, auto_add_composed_components). Attributes ---------- name : str System identifier. description : str System description. base_power : float | None System base power in MVA. See Also -------- :class:`infrasys.system.System` : Parent class with core system functionality. :class:`BaseParser` : Parser framework for building systems. """ def __init__( self, system_base: float | None = None, *, name: str | None = None, **kwargs: Any, ) -> None: """Initialize R2X Core System. This method defines the 'system_base' unit in the global Pint registry. If you create multiple System instances, the last one's system_base will be used for all unit conversions. Existing components will detect the change and issue a warning if they access system_base conversions. Parameters ---------- base_power : float, optional (defaults: 100.0) System base power in MVA for per-unit calculations. Can be provided as first positional argument or as keyword argument. name : str, optional Name of the system. If not provided, a default name will be assigned. **kwargs Additional keyword arguments passed to infrasys.System (e.g., description, auto_add_composed_components). """ merged_kwargs = dict(kwargs) if name is not None: merged_kwargs["name"] = name super_kwargs = filter_kwargs_by_signatures(merged_kwargs, callables=[InfrasysSystem]) super().__init__(**super_kwargs) self.base_power = system_base # Define the system base for pint unit conversion. # This allows components to convert: device_pu.to('system_base') units.ureg.define(f"system_base = {system_base} * MVA") # overwrite logger.debug("Setting system base to {}", system_base) def __str__(self) -> str: """Return string representation of the system. Returns ------- str String showing system name and component count. """ system_str = f"System(name={self.name}" num_components = self._components.get_num_components() if num_components: system_str += f", components={num_components}" if self.base_power: system_str += f", system_base={self.base_power}" return system_str + ")" def __repr__(self) -> str: """Return detailed string representation. Returns ------- str Same as __str__(). """ return str(self)
[docs] def add_components(self, *components: Component, **kwargs: Any) -> None: """Add one or more components to the system and set their _system_base. Parameters ---------- *components : Component Component(s) to add to the system. **kwargs Additional keyword arguments passed to parent's add_components. Notes ----- If any component is a HasPerUnit model, this method automatically sets the component's _system_base attribute for use in system-base per-unit display mode. Raises ------ ValueError If a component already has a different _system_base set. """ super().add_components(*components, **kwargs) for component in components: if isinstance(component, units.HasPerUnit): existing_base = component._get_system_base() if existing_base is not None and existing_base != self.base_power: comp_name = component.name if hasattr(component, "name") else type(component).__name__ msg = ( f"Component '{comp_name}' already has _system_base={existing_base} MVA " f"but is being added to system with base={self.base_power} MVA. " f"This may indicate the component was previously added to a different system." ) raise ValueError(msg) component._system_base = self.base_power logger.trace( "Set _system_base = {} MVA on component '{}'", self.base_power, component.name if hasattr(component, "name") else type(component).__name__, )
[docs] def to_json( # type: ignore self, fname: Path | str | None = None, overwrite: bool = False, indent: int | None = None, data: Any = None, ) -> bytes | None: """Serialize system to JSON file or return bytes. Parameters ---------- fname : Path or str, optional Output JSON file path. If None, prints JSON to stdout. Note: When writing to stdout, time series are serialized to a temporary directory that will be cleaned up automatically. overwrite : bool, default False If True, overwrite existing file. If False, raise error if file exists. indent : int, optional JSON indentation level. If None, uses compact format. data : optional Additional data to include in serialization. Returns ------- None Raises ------ FileExistsError If file exists and overwrite=False. See Also -------- :meth:`from_json` : Load system from JSON file """ if fname: return super().to_json(fname, overwrite=overwrite, indent=indent, data=data) logger.info("Serializing system '{}'", self.name) cache_folder = get_r2x_cache_path() time_series_dir = cache_folder / f"{self.uuid}_time_series" time_series_dir.mkdir(exist_ok=True, parents=True) system_data: dict[str, Any] = { "name": self.name, "description": self.description, "uuid": str(self.uuid), "data_format_version": self.data_format_version, "components": [x.model_dump_custom() for x in self._component_mgr.iter_all()], "supplemental_attributes": [ x.model_dump_custom() for x in self._supplemental_attr_mgr.iter_all() ], "time_series": { "directory": str(time_series_dir), }, } extra = self.serialize_system_attributes() system_data.update(extra) if data is None: data = system_data else: if "system" not in data: data["system"] = system_data backup(self._con, time_series_dir / self.DB_FILENAME) self._time_series_mgr.serialize(system_data["time_series"], time_series_dir, db_name=self.DB_FILENAME) json_bytes = orjson.dumps(data) return json_bytes
[docs] @classmethod def from_json( # type: ignore cls, source: Path | str | bytes, /, *, upgrade_handler: Callable[..., Any] | None = None, **kwargs: Any, ) -> "System": """Deserialize system from JSON file. Parameters ---------- source : Path, str, or bytes Input JSON source. upgrade_handler : Callable, optional Function to handle data model version upgrades. **kwargs Additional keyword arguments passed to infrasys deserialization. Returns ------- System Deserialized system instance. Raises ------ FileNotFoundError If file does not exist. ValueError If JSON format is invalid. See Also -------- :meth:`to_json` : Serialize system to JSON file. :func:`upgrade_data` : Phase 1 upgrades for parser workflow. """ match source: case Path() | str(): system = super().from_json(source, upgrade_handler=upgrade_handler, **kwargs) case bytes(): logger.debug("Deserializing system from bytes.") json_data = orjson.loads(source.decode("utf-8")) ts_info = json_data.get("time_series") if not ts_info: msg = "Data is missing time series information. Check source." raise KeyError(msg) if "directory" not in ts_info: msg = "Data is missing time series directory." raise KeyError(msg) system = super().from_dict( json_data, ts_info["directory"], upgrade_handler=upgrade_handler, **kwargs ) case _: msg = f"{type(source)=} for function from_json. Valid types are: Path, str, bytes" raise NotImplementedError(msg) for component in system.get_components(Component): if isinstance(component, units.HasPerUnit): # NOTE: mypy does not know that we deserialize the system attributes. component._system_base = system.base_power # type:ignore return system # type: ignore
[docs] def serialize_system_attributes(self) -> dict[str, Any]: """Serialize R2X-specific system attributes. Returns ------- dict[str, Any] Dictionary containing system_base_power. """ return {"system_base_power": self.base_power, "r2x_core_version": version("r2x_core")}
[docs] def deserialize_system_attributes(self, data: dict[str, Any]) -> None: """Deserialize R2X-specific system attributes. Parameters ---------- data : dict[str, Any] Dictionary containing serialized system attributes. """ if "system_base_power" in data: self.base_power = data["system_base_power"]