flixopt 2.2.0b0__py3-none-any.whl → 3.0.0__py3-none-any.whl

flixopt/io.py

@@ -1,56 +1,22 @@
+from __future__ import annotations
+
 import importlib.util
 import json
 import logging
 import pathlib
 import re
 from dataclasses import dataclass
-from typing import Dict, Literal, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Literal
 
-import linopy
 import xarray as xr
 import yaml
 
-from .core import TimeSeries
+if TYPE_CHECKING:
+    import linopy
 
 logger = logging.getLogger('flixopt')
 
 
-def replace_timeseries(obj, mode: Literal['name', 'stats', 'data'] = 'name'):
-    """Recursively replaces TimeSeries objects with their names prefixed by '::::'."""
-    if isinstance(obj, dict):
-        return {k: replace_timeseries(v, mode) for k, v in obj.items()}
-    elif isinstance(obj, list):
-        return [replace_timeseries(v, mode) for v in obj]
-    elif isinstance(obj, TimeSeries):  # Adjust this based on the actual class
-        if obj.all_equal:
-            return obj.selected_data.values.max().item()
-        elif mode == 'name':
-            return f'::::{obj.name}'
-        elif mode == 'stats':
-            return obj.stats
-        elif mode == 'data':
-            return obj
-        else:
-            raise ValueError(f'Invalid mode {mode}')
-    else:
-        return obj
-
-
-def insert_dataarray(obj, ds: xr.Dataset):
-    """Recursively inserts TimeSeries objects into a dataset."""
-    if isinstance(obj, dict):
-        return {k: insert_dataarray(v, ds) for k, v in obj.items()}
-    elif isinstance(obj, list):
-        return [insert_dataarray(v, ds) for v in obj]
-    elif isinstance(obj, str) and obj.startswith('::::'):
-        da = ds[obj[4:]]
-        if 'time' in da.dims and da.isel(time=-1).isnull().any().item():
-            return da.isel(time=slice(0, -1))
-        return da
-    else:
-        return obj
-
-
 def remove_none_and_empty(obj):
     """Recursively removes None and empty dicts and lists values from a dictionary or list."""
 
@@ -102,7 +68,7 @@ def _save_to_yaml(data, output_file='formatted_output.yaml'):
     # Configure dumper options for better formatting
     class CustomDumper(yaml.SafeDumper):
         def increase_indent(self, flow=False, indentless=False):
-            return super(CustomDumper, self).increase_indent(flow, False)
+            return super().increase_indent(flow, False)
 
     # Write to file with settings that ensure proper formatting
     with open(output_file, 'w', encoding='utf-8') as file:
@@ -175,7 +141,7 @@ def _normalize_string_content(text):
     return text.strip()
 
 
-def document_linopy_model(model: linopy.Model, path: pathlib.Path = None) -> Dict[str, str]:
+def document_linopy_model(model: linopy.Model, path: pathlib.Path | None = None) -> dict[str, str]:
     """
     Convert all model variables and constraints to a structured string representation.
     This can take multiple seconds for large models.
@@ -224,18 +190,19 @@ def document_linopy_model(model: linopy.Model, path: pathlib.Path = None) -> Dic
     if path is not None:
         if path.suffix not in ['.yaml', '.yml']:
             raise ValueError(f'Invalid file extension for path {path}. Only .yaml and .yml are supported')
-        _save_to_yaml(documentation, path)
+        _save_to_yaml(documentation, str(path))
 
     return documentation
 
 
 def save_dataset_to_netcdf(
     ds: xr.Dataset,
-    path: Union[str, pathlib.Path],
+    path: str | pathlib.Path,
     compression: int = 0,
+    engine: Literal['netcdf4', 'scipy', 'h5netcdf'] = 'h5netcdf',
 ) -> None:
     """
-    Save a dataset to a netcdf file. Store the attrs as a json string in the 'attrs' attribute.
+    Save a dataset to a netcdf file. Store all attrs as JSON strings in 'attrs' attributes.
 
     Args:
         ds: Dataset to save.
@@ -245,40 +212,68 @@ def save_dataset_to_netcdf(
     Raises:
         ValueError: If the path has an invalid file extension.
     """
+    path = pathlib.Path(path)
     if path.suffix not in ['.nc', '.nc4']:
         raise ValueError(f'Invalid file extension for path {path}. Only .nc and .nc4 are supported')
 
     apply_encoding = False
     if compression != 0:
-        if importlib.util.find_spec('netCDF4') is not None:
+        if importlib.util.find_spec(engine) is not None:
             apply_encoding = True
         else:
             logger.warning(
-                'Dataset was exported without compression due to missing dependency "netcdf4".'
-                'Install netcdf4 via `pip install netcdf4`.'
+                f'Dataset was exported without compression due to missing dependency "{engine}".'
+                f'Install {engine} via `pip install {engine}`.'
             )
+
     ds = ds.copy(deep=True)
     ds.attrs = {'attrs': json.dumps(ds.attrs)}
+
+    # Convert all DataArray attrs to JSON strings
+    for var_name, data_var in ds.data_vars.items():
+        if data_var.attrs:  # Only if there are attrs
+            ds[var_name].attrs = {'attrs': json.dumps(data_var.attrs)}
+
+    # Also handle coordinate attrs if they exist
+    for coord_name, coord_var in ds.coords.items():
+        if hasattr(coord_var, 'attrs') and coord_var.attrs:
+            ds[coord_name].attrs = {'attrs': json.dumps(coord_var.attrs)}
+
     ds.to_netcdf(
         path,
         encoding=None
         if not apply_encoding
-        else {data_var: {'zlib': True, 'complevel': 5} for data_var in ds.data_vars},
+        else {data_var: {'zlib': True, 'complevel': compression} for data_var in ds.data_vars},
+        engine=engine,
     )
 
 
-def load_dataset_from_netcdf(path: Union[str, pathlib.Path]) -> xr.Dataset:
+def load_dataset_from_netcdf(path: str | pathlib.Path) -> xr.Dataset:
     """
-    Load a dataset from a netcdf file. Load the attrs from the 'attrs' attribute.
+    Load a dataset from a netcdf file. Load all attrs from 'attrs' attributes.
 
     Args:
         path: Path to load the dataset from.
 
     Returns:
-        Dataset: Loaded dataset.
+        Dataset: Loaded dataset with restored attrs.
     """
-    ds = xr.load_dataset(path)
-    ds.attrs = json.loads(ds.attrs['attrs'])
+    ds = xr.load_dataset(str(path), engine='h5netcdf')
+
+    # Restore Dataset attrs
+    if 'attrs' in ds.attrs:
+        ds.attrs = json.loads(ds.attrs['attrs'])
+
+    # Restore DataArray attrs
+    for var_name, data_var in ds.data_vars.items():
+        if 'attrs' in data_var.attrs:
+            ds[var_name].attrs = json.loads(data_var.attrs['attrs'])
+
+    # Restore coordinate attrs
+    for coord_name, coord_var in ds.coords.items():
+        if hasattr(coord_var, 'attrs') and 'attrs' in coord_var.attrs:
+            ds[coord_name].attrs = json.loads(coord_var.attrs['attrs'])
+
     return ds
 
 
@@ -302,7 +297,7 @@ class CalculationResultsPaths:
         self.flow_system = self.folder / f'{self.name}--flow_system.nc4'
         self.model_documentation = self.folder / f'{self.name}--model_documentation.yaml'
 
-    def all_paths(self) -> Dict[str, pathlib.Path]:
+    def all_paths(self) -> dict[str, pathlib.Path]:
         """Return a dictionary of all paths."""
         return {
             'linopy_model': self.linopy_model,
@@ -326,7 +321,7 @@ class CalculationResultsPaths:
                     f'Folder {self.folder} and its parent do not exist. Please create them first.'
                 ) from e
 
-    def update(self, new_name: Optional[str] = None, new_folder: Optional[pathlib.Path] = None) -> None:
+    def update(self, new_name: str | None = None, new_folder: pathlib.Path | None = None) -> None:
         """Update name and/or folder and refresh all paths."""
         if new_name is not None:
             self.name = new_name