flixopt 3.0.1__py3-none-any.whl → 6.0.0rc7__py3-none-any.whl

flixopt/structure.py

@@ -8,34 +8,140 @@ from __future__ import annotations
 import inspect
 import json
 import logging
+import pathlib
+import re
+import warnings
 from dataclasses import dataclass
-from io import StringIO
+from difflib import get_close_matches
+from enum import Enum
 from typing import (
     TYPE_CHECKING,
     Any,
+    ClassVar,
+    Generic,
     Literal,
+    TypeVar,
 )
 
 import linopy
 import numpy as np
 import pandas as pd
 import xarray as xr
-from rich.console import Console
-from rich.pretty import Pretty
 
 from . import io as fx_io
-from .core import TimeSeriesData, get_dataarray_stats
+from .config import DEPRECATION_REMOVAL_VERSION
+from .core import FlowSystemDimensions, TimeSeriesData, get_dataarray_stats
 
 if TYPE_CHECKING:  # for type checking and preventing circular imports
-    import pathlib
     from collections.abc import Collection, ItemsView, Iterator
 
     from .effects import EffectCollectionModel
     from .flow_system import FlowSystem
+    from .types import Effect_TPS, Numeric_TPS, NumericOrBool
 
 logger = logging.getLogger('flixopt')
 
 
+def _ensure_coords(
+    data: xr.DataArray | float | int,
+    coords: xr.Coordinates | dict,
+) -> xr.DataArray | float:
+    """Broadcast data to coords if needed.
+
+    This is used at the linopy interface to ensure bounds are properly broadcasted
+    to the target variable shape. Linopy needs at least one bound to have all
+    dimensions to determine the variable shape.
+
+    Note: Infinity values (-inf, inf) are kept as scalars because linopy uses
+    special checks like `if (lower != -inf)` that fail with DataArrays.
+    """
+    # Handle both dict and xr.Coordinates
+    if isinstance(coords, dict):
+        coord_dims = list(coords.keys())
+    else:
+        coord_dims = list(coords.dims)
+
+    # Keep infinity values as scalars (linopy uses them for special checks)
+    if not isinstance(data, xr.DataArray):
+        if np.isinf(data):
+            return data
+        # Finite scalar - create full DataArray
+        return xr.DataArray(data, coords=coords, dims=coord_dims)
+
+    if set(data.dims) == set(coord_dims):
+        # Has all dims - ensure correct order
+        if data.dims != tuple(coord_dims):
+            return data.transpose(*coord_dims)
+        return data
+
+    # Broadcast to full coords (broadcast_like ensures correct dim order)
+    template = xr.DataArray(coords=coords, dims=coord_dims)
+    return data.broadcast_like(template)
+
+
+class VariableCategory(Enum):
+    """Fine-grained variable categories - names mirror variable names.
+
+    Each variable type has its own category for precise handling during
+    segment expansion and statistics calculation.
+    """
+
+    # === State variables ===
+    CHARGE_STATE = 'charge_state'  # Storage SOC (interpolate between boundaries)
+    SOC_BOUNDARY = 'soc_boundary'  # Intercluster SOC boundaries
+
+    # === Rate/Power variables ===
+    FLOW_RATE = 'flow_rate'  # Flow rate (kW)
+    NETTO_DISCHARGE = 'netto_discharge'  # Storage net discharge
+    VIRTUAL_FLOW = 'virtual_flow'  # Bus penalty slack variables
+
+    # === Binary state ===
+    STATUS = 'status'  # On/off status (persists through segment)
+    INACTIVE = 'inactive'  # Complementary inactive status
+
+    # === Binary events ===
+    STARTUP = 'startup'  # Startup event
+    SHUTDOWN = 'shutdown'  # Shutdown event
+
+    # === Effect variables ===
+    PER_TIMESTEP = 'per_timestep'  # Effect per timestep
+    SHARE = 'share'  # All temporal contributions (flow, active, startup)
+    TOTAL = 'total'  # Effect total (per period/scenario)
+    TOTAL_OVER_PERIODS = 'total_over_periods'  # Effect total over all periods
+
+    # === Investment ===
+    SIZE = 'size'  # Generic investment size (for backwards compatibility)
+    FLOW_SIZE = 'flow_size'  # Flow investment size
+    STORAGE_SIZE = 'storage_size'  # Storage capacity size
+    INVESTED = 'invested'  # Invested yes/no binary
+
+    # === Counting/Duration ===
+    STARTUP_COUNT = 'startup_count'  # Count of startups
+    DURATION = 'duration'  # Duration tracking (uptime/downtime)
+
+    # === Piecewise linearization ===
+    INSIDE_PIECE = 'inside_piece'  # Binary segment selection
+    LAMBDA0 = 'lambda0'  # Interpolation weight
+    LAMBDA1 = 'lambda1'  # Interpolation weight
+    ZERO_POINT = 'zero_point'  # Zero point handling
+
+    # === Other ===
+    OTHER = 'other'  # Uncategorized
+
+
+# === Logical Groupings for Segment Expansion ===
+# Default behavior (not listed): repeat value within segment
+
+EXPAND_INTERPOLATE: set[VariableCategory] = {VariableCategory.CHARGE_STATE}
+"""State variables that should be interpolated between segment boundaries."""
+
+EXPAND_DIVIDE: set[VariableCategory] = {VariableCategory.PER_TIMESTEP, VariableCategory.SHARE}
+"""Segment totals that should be divided by expansion factor to preserve sums."""
+
+EXPAND_FIRST_TIMESTEP: set[VariableCategory] = {VariableCategory.STARTUP, VariableCategory.SHUTDOWN}
+"""Binary events that should appear only at the first timestep of the segment."""
+
+
 CLASS_REGISTRY = {}
 
 
@@ -86,17 +192,35 @@ class FlowSystemModel(linopy.Model, SubmodelsMixin):
 
     Args:
         flow_system: The flow_system that is used to create the model.
-        normalize_weights: Whether to automatically normalize the weights to sum up to 1 when solving.
     """
 
-    def __init__(self, flow_system: FlowSystem, normalize_weights: bool):
+    def __init__(self, flow_system: FlowSystem):
         super().__init__(force_dim_names=True)
         self.flow_system = flow_system
-        self.normalize_weights = normalize_weights
         self.effects: EffectCollectionModel | None = None
         self.submodels: Submodels = Submodels({})
+        self.variable_categories: dict[str, VariableCategory] = {}
+
+    def add_variables(
+        self,
+        lower: xr.DataArray | float = -np.inf,
+        upper: xr.DataArray | float = np.inf,
+        coords: xr.Coordinates | None = None,
+        **kwargs,
+    ) -> linopy.Variable:
+        """Override to ensure bounds are broadcasted to coords shape.
+
+        Linopy uses the union of all DataArray dimensions to determine variable shape.
+        This override ensures at least one bound has all target dimensions when coords
+        is provided, allowing internal data to remain compact (scalars, 1D arrays).
+        """
+        if coords is not None:
+            lower = _ensure_coords(lower, coords)
+            upper = _ensure_coords(upper, coords)
+        return super().add_variables(lower=lower, upper=upper, coords=coords, **kwargs)
 
     def do_modeling(self):
+        # Create all element models
         self.effects = self.flow_system.effects.create_model(self)
         for component in self.flow_system.components.values():
             component.create_model(self)
@@ -106,6 +230,16 @@ class FlowSystemModel(linopy.Model, SubmodelsMixin):
         # Add scenario equality constraints after all elements are modeled
         self._add_scenario_equality_constraints()
 
+        # Populate _variable_names and _constraint_names on each Element
+        self._populate_element_variable_names()
+
+    def _populate_element_variable_names(self):
+        """Populate _variable_names and _constraint_names on each Element from its submodel."""
+        for element in self.flow_system.values():
+            if element.submodel is not None:
+                element._variable_names = list(element.submodel.variables)
+                element._constraint_names = list(element.submodel.constraints)
+
     def _add_scenario_equality_for_parameter_type(
         self,
         parameter_type: Literal['flow_rate', 'size'],
@@ -154,38 +288,130 @@ class FlowSystemModel(linopy.Model, SubmodelsMixin):
 
     @property
     def solution(self):
-        solution = super().solution
+        """Build solution dataset, reindexing to timesteps_extra for consistency."""
+        # Suppress the linopy warning about coordinate mismatch.
+        # This warning is expected when storage charge_state has one more timestep than other variables.
+        with warnings.catch_warnings():
+            warnings.filterwarnings(
+                'ignore',
+                category=UserWarning,
+                message='Coordinates across variables not equal',
+            )
+            solution = super().solution
         solution['objective'] = self.objective.value
+        # Store attrs as JSON strings for netCDF compatibility
         solution.attrs = {
-            'Components': {
-                comp.label_full: comp.submodel.results_structure()
-                for comp in sorted(
-                    self.flow_system.components.values(), key=lambda component: component.label_full.upper()
-                )
-            },
-            'Buses': {
-                bus.label_full: bus.submodel.results_structure()
-                for bus in sorted(self.flow_system.buses.values(), key=lambda bus: bus.label_full.upper())
-            },
-            'Effects': {
-                effect.label_full: effect.submodel.results_structure()
-                for effect in sorted(self.flow_system.effects, key=lambda effect: effect.label_full.upper())
-            },
-            'Flows': {
-                flow.label_full: flow.submodel.results_structure()
-                for flow in sorted(self.flow_system.flows.values(), key=lambda flow: flow.label_full.upper())
-            },
+            'Components': json.dumps(
+                {
+                    comp.label_full: comp.submodel.results_structure()
+                    for comp in sorted(
+                        self.flow_system.components.values(), key=lambda component: component.label_full.upper()
+                    )
+                }
+            ),
+            'Buses': json.dumps(
+                {
+                    bus.label_full: bus.submodel.results_structure()
+                    for bus in sorted(self.flow_system.buses.values(), key=lambda bus: bus.label_full.upper())
+                }
+            ),
+            'Effects': json.dumps(
+                {
+                    effect.label_full: effect.submodel.results_structure()
+                    for effect in sorted(
+                        self.flow_system.effects.values(), key=lambda effect: effect.label_full.upper()
+                    )
+                }
+            ),
+            'Flows': json.dumps(
+                {
+                    flow.label_full: flow.submodel.results_structure()
+                    for flow in sorted(self.flow_system.flows.values(), key=lambda flow: flow.label_full.upper())
+                }
+            ),
         }
-        return solution.reindex(time=self.flow_system.timesteps_extra)
+        # Ensure solution is always indexed by timesteps_extra for consistency.
+        # Variables without extra timestep data will have NaN at the final timestep.
+        if 'time' in solution.coords:
+            if not solution.indexes['time'].equals(self.flow_system.timesteps_extra):
+                solution = solution.reindex(time=self.flow_system.timesteps_extra)
+        return solution
 
     @property
-    def hours_per_step(self):
-        return self.flow_system.hours_per_timestep
+    def timestep_duration(self) -> xr.DataArray:
+        """Duration of each timestep in hours."""
+        return self.flow_system.timestep_duration
 
     @property
     def hours_of_previous_timesteps(self):
         return self.flow_system.hours_of_previous_timesteps
 
+    @property
+    def dims(self) -> list[str]:
+        """Active dimension names."""
+        return self.flow_system.dims
+
+    @property
+    def indexes(self) -> dict[str, pd.Index]:
+        """Indexes for active dimensions."""
+        return self.flow_system.indexes
+
+    @property
+    def weights(self) -> dict[str, xr.DataArray]:
+        """Weights for active dimensions (unit weights if not set).
+
+        Scenario weights are always normalized (handled by FlowSystem).
+        """
+        return self.flow_system.weights
+
+    @property
+    def temporal_dims(self) -> list[str]:
+        """Temporal dimensions for summing over time.
+
+        Returns ['time', 'cluster'] for clustered systems, ['time'] otherwise.
+        """
+        return self.flow_system.temporal_dims
+
+    @property
+    def temporal_weight(self) -> xr.DataArray:
+        """Combined temporal weight (timestep_duration × cluster_weight)."""
+        return self.flow_system.temporal_weight
+
+    def sum_temporal(self, data: xr.DataArray) -> xr.DataArray:
+        """Sum data over temporal dimensions with full temporal weighting.
+
+        Example:
+            >>> total_energy = model.sum_temporal(flow_rate)
+        """
+        return self.flow_system.sum_temporal(data)
+
+    @property
+    def scenario_weights(self) -> xr.DataArray:
+        """Scenario weights of model.
+
+        Returns:
+            - Scalar 1 if no scenarios defined
+            - Unit weights (all 1.0) if scenarios exist but no explicit weights set
+            - Normalized explicit weights if set via FlowSystem.scenario_weights
+        """
+        if self.flow_system.scenarios is None:
+            return xr.DataArray(1)
+
+        if self.flow_system.scenario_weights is None:
+            return self.flow_system._unit_weight('scenario')
+
+        return self.flow_system.scenario_weights
+
+    @property
+    def objective_weights(self) -> xr.DataArray:
+        """
+        Objective weights of model (period_weights × scenario_weights).
+        """
+        period_weights = self.flow_system.effects.objective_effect.submodel.period_weights
+        scenario_weights = self.scenario_weights
+
+        return period_weights * scenario_weights
+
     def get_coords(
         self,
         dims: Collection[str] | None = None,
@@ -196,7 +422,8 @@ class FlowSystemModel(linopy.Model, SubmodelsMixin):
 
         Args:
             dims: The dimensions to include in the coordinates. If None, includes all dimensions
-            extra_timestep: If True, uses extra timesteps instead of regular timesteps
+            extra_timestep: If True, uses extra timesteps instead of regular timesteps.
+                For clustered FlowSystems, extends time by 1 (for charge_state boundaries).
 
         Returns:
             The coordinates of the model, or None if no coordinates are available
@@ -208,28 +435,20 @@ class FlowSystemModel(linopy.Model, SubmodelsMixin):
             raise ValueError('extra_timestep=True requires "time" to be included in dims')
 
         if dims is None:
-            coords = dict(self.flow_system.coords)
+            coords = dict(self.flow_system.indexes)
         else:
-            coords = {k: v for k, v in self.flow_system.coords.items() if k in dims}
+            # In clustered systems, 'time' is always paired with 'cluster'
+            # So when 'time' is requested, also include 'cluster' if available
+            effective_dims = set(dims)
+            if 'time' in dims and 'cluster' in self.flow_system.indexes:
+                effective_dims.add('cluster')
+            coords = {k: v for k, v in self.flow_system.indexes.items() if k in effective_dims}
 
         if extra_timestep and coords:
             coords['time'] = self.flow_system.timesteps_extra
 
         return xr.Coordinates(coords) if coords else None
 
-    @property
-    def weights(self) -> int | xr.DataArray:
-        """Returns the weights of the FlowSystem. Normalizes to 1 if normalize_weights is True"""
-        if self.flow_system.weights is not None:
-            weights = self.flow_system.weights
-        else:
-            weights = self.flow_system.fit_to_model_coords('weights', 1, dims=['period', 'scenario'])
-
-        if not self.normalize_weights:
-            return weights
-
-        return weights / weights.sum()
-
     def __repr__(self) -> str:
         """
         Return a string representation of the FlowSystemModel, borrowed from linopy.Model.
@@ -243,9 +462,7 @@ class FlowSystemModel(linopy.Model, SubmodelsMixin):
         }
 
         # Format sections with headers and underlines
-        formatted_sections = []
-        for section_header, section_content in sections.items():
-            formatted_sections.append(f'{section_header}\n{"-" * len(section_header)}\n{section_content}')
+        formatted_sections = fx_io.format_sections_with_headers(sections)
 
         title = f'FlowSystemModel ({self.type})'
         all_sections = '\n'.join(formatted_sections)
@@ -268,21 +485,133 @@ class Interface:
         - Recursive handling of complex nested structures
 
     Subclasses must implement:
-        transform_data(flow_system): Transform data to match FlowSystem dimensions
+        transform_data(): Transform data to match FlowSystem dimensions
     """
 
-    def transform_data(self, flow_system: FlowSystem, name_prefix: str = '') -> None:
+    # Class-level defaults for attributes set by link_to_flow_system()
+    # These provide type hints and default values without requiring __init__ in subclasses
+    _flow_system: FlowSystem | None = None
+    _prefix: str = ''
+
+    def transform_data(self) -> None:
         """Transform the data of the interface to match the FlowSystem's dimensions.
 
-        Args:
-            flow_system: The FlowSystem containing timing and dimensional information
-            name_prefix: The prefix to use for the names of the variables. Defaults to '', which results in no prefix.
+        Uses `self._prefix` (set during `link_to_flow_system()`) to name transformed data.
 
         Raises:
             NotImplementedError: Must be implemented by subclasses
+
+        Note:
+            The FlowSystem reference is available via self._flow_system (for Interface objects)
+            or self.flow_system property (for Element objects). Elements must be registered
+            to a FlowSystem before calling this method.
         """
         raise NotImplementedError('Every Interface subclass needs a transform_data() method')
 
+    @property
+    def prefix(self) -> str:
+        """The prefix used for naming transformed data (e.g., 'Boiler(Q_th)|status_parameters')."""
+        return self._prefix
+
+    def _sub_prefix(self, name: str) -> str:
+        """Build a prefix for a nested interface by appending name to current prefix."""
+        return f'{self._prefix}|{name}' if self._prefix else name
+
+    def link_to_flow_system(self, flow_system: FlowSystem, prefix: str = '') -> None:
+        """Link this interface and all nested interfaces to a FlowSystem.
+
+        This method is called automatically during element registration to enable
+        elements to access FlowSystem properties without passing the reference
+        through every method call. It also sets the prefix used for naming
+        transformed data.
+
+        Subclasses with nested Interface objects should override this method
+        to propagate the link to their nested interfaces by calling
+        `super().link_to_flow_system(flow_system, prefix)` first, then linking
+        nested objects with appropriate prefixes.
+
+        Args:
+            flow_system: The FlowSystem to link to
+            prefix: The prefix for naming transformed data (e.g., 'Boiler(Q_th)')
+
+        Examples:
+            Override in a subclass with nested interfaces:
+
+            ```python
+            def link_to_flow_system(self, flow_system, prefix: str = '') -> None:
+                super().link_to_flow_system(flow_system, prefix)
+                if self.nested_interface is not None:
+                    self.nested_interface.link_to_flow_system(flow_system, f'{prefix}|nested' if prefix else 'nested')
+            ```
+
+            Creating an Interface dynamically during modeling:
+
+            ```python
+            # In a Model class
+            if flow.status_parameters is None:
+                flow.status_parameters = StatusParameters()
+                flow.status_parameters.link_to_flow_system(self._model.flow_system, f'{flow.label_full}')
+            ```
+        """
+        self._flow_system = flow_system
+        self._prefix = prefix
+
+    @property
+    def flow_system(self) -> FlowSystem:
+        """Access the FlowSystem this interface is linked to.
+
+        Returns:
+            The FlowSystem instance this interface belongs to.
+
+        Raises:
+            RuntimeError: If interface has not been linked to a FlowSystem yet.
+
+        Note:
+            For Elements, this is set during add_elements().
+            For parameter classes, this is set recursively when the parent Element is registered.
+        """
+        if self._flow_system is None:
+            raise RuntimeError(
+                f'{self.__class__.__name__} is not linked to a FlowSystem. '
+                f'Ensure the parent element is registered via flow_system.add_elements() first.'
+            )
+        return self._flow_system
+
+    def _fit_coords(
+        self, name: str, data: NumericOrBool | None, dims: Collection[FlowSystemDimensions] | None = None
+    ) -> xr.DataArray | None:
+        """Convenience wrapper for FlowSystem.fit_to_model_coords().
+
+        Args:
+            name: The name for the data variable
+            data: The data to transform
+            dims: Optional dimension names
+
+        Returns:
+            Transformed data aligned to FlowSystem coordinates
+        """
+        return self.flow_system.fit_to_model_coords(name, data, dims=dims)
+
+    def _fit_effect_coords(
+        self,
+        prefix: str | None,
+        effect_values: Effect_TPS | Numeric_TPS | None,
+        suffix: str | None = None,
+        dims: Collection[FlowSystemDimensions] | None = None,
+    ) -> Effect_TPS | None:
+        """Convenience wrapper for FlowSystem.fit_effects_to_model_coords().
+
+        Args:
+            prefix: Label prefix for effect names
+            effect_values: The effect values to transform
+            suffix: Optional label suffix
+            dims: Optional dimension names
+
+        Returns:
+            Transformed effect values aligned to FlowSystem coordinates
+        """
+        return self.flow_system.fit_effects_to_model_coords(prefix, effect_values, suffix, dims=dims)
+
     def _create_reference_structure(self) -> tuple[dict, dict[str, xr.DataArray]]:
         """
         Convert all DataArrays to references and extract them.
@@ -424,6 +753,7 @@ class Interface:
         current_value: Any = None,
         transform: callable = None,
         check_conflict: bool = True,
+        additional_warning_message: str = '',
     ) -> Any:
         """
         Handle a deprecated keyword argument by issuing a warning and returning the appropriate value.
@@ -439,6 +769,7 @@ class Interface:
             check_conflict: Whether to check if both old and new parameters are specified (default: True).
                 Note: For parameters with non-None default values (e.g., bool parameters with default=False),
                 set check_conflict=False since we cannot distinguish between an explicit value and the default.
+            additional_warning_message: Add a custom message which gets appended with a line break to the default warning.
 
         Returns:
             The value to use (either from old parameter or current_value)
@@ -461,8 +792,18 @@ class Interface:
 
         old_value = kwargs.pop(old_name, None)
         if old_value is not None:
+            # Build base warning message
+            base_warning = f'The use of the "{old_name}" argument is deprecated. Use the "{new_name}" argument instead. Will be removed in v{DEPRECATION_REMOVAL_VERSION}.'
+
+            # Append additional message on a new line if provided
+            if additional_warning_message:
+                # Normalize whitespace: strip leading/trailing whitespace
+                extra_msg = additional_warning_message.strip()
+                if extra_msg:
+                    base_warning += '\n' + extra_msg
+
             warnings.warn(
-                f'The use of the "{old_name}" argument is deprecated. Use the "{new_name}" argument instead.',
+                base_warning,
                 DeprecationWarning,
                 stacklevel=3,  # Stack: this method -> __init__ -> caller
             )
@@ -507,6 +848,33 @@ class Interface:
             unexpected_params = ', '.join(f"'{param}'" for param in extra_kwargs.keys())
             raise TypeError(f'{class_name}.__init__() got unexpected keyword argument(s): {unexpected_params}')
 
+    @staticmethod
+    def _has_value(param: Any) -> bool:
+        """Check if a parameter has a meaningful value.
+
+        Args:
+            param: The parameter to check.
+
+        Returns:
+            False for:
+                - None
+                - Empty collections (dict, list, tuple, set, frozenset)
+
+            True for all other values, including:
+                - Non-empty collections
+                - xarray DataArrays (even if they contain NaN/empty data)
+                - Scalar values (0, False, empty strings, etc.)
+                - NumPy arrays (even if empty - use .size to check those explicitly)
+        """
+        if param is None:
+            return False
+
+        # Check for empty collections (but not strings, arrays, or DataArrays)
+        if isinstance(param, (dict, list, tuple, set, frozenset)) and len(param) == 0:
+            return False
+
+        return True
+
     @classmethod
     def _resolve_dataarray_reference(
         cls, reference: str, arrays_dict: dict[str, xr.DataArray]
@@ -530,8 +898,11 @@ class Interface:
 
         array = arrays_dict[array_name]
 
-        # Handle null values with warning
-        if array.isnull().any():
+        # Handle null values with warning (use numpy for performance - 200x faster than xarray)
+        has_nulls = (np.issubdtype(array.dtype, np.floating) and np.any(np.isnan(array.values))) or (
+            array.dtype == object and pd.isna(array.values).any()
+        )
+        if has_nulls:
             logger.error(f"DataArray '{array_name}' contains null values. Dropping all-null along present dims.")
             if 'time' in array.dims:
                 array = array.dropna(dim='time', how='all')
@@ -586,7 +957,34 @@ class Interface:
                 resolved_nested_data = cls._resolve_reference_structure(nested_data, arrays_dict)
 
                 try:
-                    return nested_class(**resolved_nested_data)
+                    # Get valid constructor parameters for this class
+                    init_params = set(inspect.signature(nested_class.__init__).parameters.keys())
+
+                    # Check for deferred init attributes (defined as class attribute on Element subclasses)
+                    # These are serialized but set after construction, not passed to child __init__
+                    deferred_attr_names = getattr(nested_class, '_deferred_init_attrs', set())
+                    deferred_attrs = {k: v for k, v in resolved_nested_data.items() if k in deferred_attr_names}
+                    constructor_data = {k: v for k, v in resolved_nested_data.items() if k not in deferred_attr_names}
+
+                    # Check for unknown parameters - these could be typos or renamed params
+                    unknown_params = set(constructor_data.keys()) - init_params
+                    if unknown_params:
+                        raise TypeError(
+                            f'{class_name}.__init__() got unexpected keyword arguments: {unknown_params}. '
+                            f'This may indicate renamed parameters that need conversion. '
+                            f'Valid parameters are: {init_params - {"self"}}'
+                        )
+
+                    # Create instance with constructor parameters
+                    instance = nested_class(**constructor_data)
+
+                    # Set internal attributes after construction
+                    for attr_name, attr_value in deferred_attrs.items():
+                        setattr(instance, attr_name, attr_value)
+
+                    return instance
+                except TypeError as e:
+                    raise ValueError(f'Failed to create instance of {class_name}: {e}') from e
                 except Exception as e:
                     raise ValueError(f'Failed to create instance of {class_name}: {e}') from e
             else:
@@ -662,18 +1060,29 @@ class Interface:
                 f'Original Error: {e}'
             ) from e
 
-    def to_netcdf(self, path: str | pathlib.Path, compression: int = 0):
+    def to_netcdf(self, path: str | pathlib.Path, compression: int = 5, overwrite: bool = False):
         """
         Save the object to a NetCDF file.
 
         Args:
-            path: Path to save the NetCDF file
+            path: Path to save the NetCDF file. Parent directories are created if they don't exist.
             compression: Compression level (0-9)
+            overwrite: If True, overwrite existing file. If False, raise error if file exists.
 
         Raises:
+            FileExistsError: If overwrite=False and file already exists.
             ValueError: If serialization fails
             IOError: If file cannot be written
         """
+        path = pathlib.Path(path)
+
+        # Check if file exists (unless overwrite is True)
+        if not overwrite and path.exists():
+            raise FileExistsError(f'File already exists: {path}. Use overwrite=True to overwrite existing file.')
+
+        # Create parent directories if they don't exist
+        path.parent.mkdir(parents=True, exist_ok=True)
+
         try:
             ds = self.to_dataset()
             fx_io.save_dataset_to_netcdf(ds, path, compression=compression)
@@ -707,7 +1116,19 @@ class Interface:
             reference_structure.pop('__class__', None)
 
             # Create arrays dictionary from dataset variables
-            arrays_dict = {name: array for name, array in ds.data_vars.items()}
+            # Use ds.variables with coord_cache for faster DataArray construction
+            variables = ds.variables
+            coord_cache = {k: ds.coords[k] for k in ds.coords}
+            coord_names = set(coord_cache)
+            arrays_dict = {
+                name: xr.DataArray(
+                    variables[name],
+                    coords={k: coord_cache[k] for k in variables[name].dims if k in coord_cache},
+                    name=name,
+                )
+                for name in variables
+                if name not in coord_names
+            }
 
             # Resolve all references using the centralized method
             resolved_params = cls._resolve_reference_structure(reference_structure, arrays_dict)
@@ -788,47 +1209,13 @@ class Interface:
         try:
             # Use the stats mode for JSON export (cleaner output)
             data = self.get_structure(clean=True, stats=True)
-            with open(path, 'w', encoding='utf-8') as f:
-                json.dump(data, f, indent=4, ensure_ascii=False)
+            fx_io.save_json(data, path)
         except Exception as e:
             raise OSError(f'Failed to save {self.__class__.__name__} to JSON file {path}: {e}') from e
 
     def __repr__(self):
         """Return a detailed string representation for debugging."""
-        try:
-            # Get the constructor arguments and their current values
-            init_signature = inspect.signature(self.__init__)
-            init_args = init_signature.parameters
-
-            # Create a dictionary with argument names and their values, with better formatting
-            args_parts = []
-            for name in init_args:
-                if name == 'self':
-                    continue
-                value = getattr(self, name, None)
-                # Truncate long representations
-                value_repr = repr(value)
-                if len(value_repr) > 50:
-                    value_repr = value_repr[:47] + '...'
-                args_parts.append(f'{name}={value_repr}')
-
-            args_str = ', '.join(args_parts)
-            return f'{self.__class__.__name__}({args_str})'
-        except Exception:
-            # Fallback if introspection fails
-            return f'{self.__class__.__name__}(<repr_failed>)'
-
-    def __str__(self):
-        """Return a user-friendly string representation."""
-        try:
-            data = self.get_structure(clean=True, stats=True)
-            with StringIO() as output_buffer:
-                console = Console(file=output_buffer, width=1000)  # Adjust width as needed
-                console.print(Pretty(data, expand_all=True, indent_guides=True))
-                return output_buffer.getvalue()
-        except Exception:
-            # Fallback if structure generation fails
-            return f'{self.__class__.__name__} instance'
+        return fx_io.build_repr_from_init(self, excluded_params={'self', 'label', 'kwargs'})
 
     def copy(self) -> Interface:
         """
@@ -856,15 +1243,36 @@ class Interface:
 class Element(Interface):
     """This class is the basic Element of flixopt. Every Element has a label"""
 
-    def __init__(self, label: str, meta_data: dict | None = None):
+    submodel: ElementModel | None
+
+    # Attributes that are serialized but set after construction (not passed to child __init__)
+    # These are internal state populated during modeling, not user-facing parameters
+    _deferred_init_attrs: ClassVar[set[str]] = {'_variable_names', '_constraint_names'}
+
+    def __init__(
+        self,
+        label: str,
+        meta_data: dict | None = None,
+        color: str | None = None,
+        _variable_names: list[str] | None = None,
+        _constraint_names: list[str] | None = None,
+    ):
         """
         Args:
             label: The label of the element
             meta_data: used to store more information about the Element. Is not used internally, but saved in the results. Only use python native types.
+            color: Optional color for visualizations (e.g., '#FF6B6B'). If not provided, a color will be automatically assigned during FlowSystem.connect_and_transform().
+            _variable_names: Internal. Variable names for this element (populated after modeling).
+            _constraint_names: Internal. Constraint names for this element (populated after modeling).
         """
         self.label = Element._valid_label(label)
         self.meta_data = meta_data if meta_data is not None else {}
-        self.submodel: ElementModel | None = None
+        self.color = color
+        self.submodel = None
+        self._flow_system: FlowSystem | None = None
+        # Variable/constraint names - populated after modeling, serialized for results
+        self._variable_names: list[str] = _variable_names if _variable_names is not None else []
+        self._constraint_names: list[str] = _constraint_names if _constraint_names is not None else []
 
     def _plausibility_checks(self) -> None:
         """This function is used to do some basic plausibility checks for each Element during initialization.
@@ -878,15 +1286,50 @@ class Element(Interface):
     def label_full(self) -> str:
         return self.label
 
+    @property
+    def solution(self) -> xr.Dataset:
+        """Solution data for this element's variables.
+
+        Returns a view into FlowSystem.solution containing only this element's variables.
+
+        Raises:
+            ValueError: If no solution is available (optimization not run or not solved).
+        """
+        if self._flow_system is None:
+            raise ValueError(f'Element "{self.label}" is not linked to a FlowSystem.')
+        if self._flow_system.solution is None:
+            raise ValueError(f'No solution available for "{self.label}". Run optimization first or load results.')
+        if not self._variable_names:
+            raise ValueError(f'No variable names available for "{self.label}". Element may not have been modeled yet.')
+        return self._flow_system.solution[self._variable_names]
+
+    def _create_reference_structure(self) -> tuple[dict, dict[str, xr.DataArray]]:
+        """
+        Override to include _variable_names and _constraint_names in serialization.
+
+        These attributes are defined in Element but may not be in subclass constructors,
+        so we need to add them explicitly.
+        """
+        reference_structure, all_extracted_arrays = super()._create_reference_structure()
+
+        # Always include variable/constraint names for solution access after loading
+        if self._variable_names:
+            reference_structure['_variable_names'] = self._variable_names
+        if self._constraint_names:
+            reference_structure['_constraint_names'] = self._constraint_names
+
+        return reference_structure, all_extracted_arrays
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return fx_io.build_repr_from_init(self, excluded_params={'self', 'label', 'kwargs'}, skip_default_size=True)
+
     @staticmethod
     def _valid_label(label: str) -> str:
-        """
-        Checks if the label is valid. If not, it is replaced by the default label
+        """Checks if the label is valid. If not, it is replaced by the default label.
 
-        Raises
-        ------
-        ValueError
-            If the label is not valid
+        Raises:
+            ValueError: If the label is not valid.
         """
         not_allowed = ['(', ')', '|', '->', '\\', '-slash-']  # \\ is needed to check for \
         if any([sign in label for sign in not_allowed]):
@@ -900,6 +1343,378 @@ class Element(Interface):
         return label
 
 
+# Precompiled regex pattern for natural sorting
+_NATURAL_SPLIT = re.compile(r'(\d+)')
+
+
+def _natural_sort_key(text):
+    """Sort key for natural ordering (e.g., bus1, bus2, bus10 instead of bus1, bus10, bus2)."""
+    return [int(c) if c.isdigit() else c.lower() for c in _NATURAL_SPLIT.split(text)]
+
+
+# Type variable for containers
+T = TypeVar('T')
+
+
+class ContainerMixin(dict[str, T]):
+    """
+    Mixin providing shared container functionality with nice repr and error messages.
+
+    Subclasses must implement _get_label() to extract the label from elements.
+    """
+
+    def __init__(
+        self,
+        elements: list[T] | dict[str, T] | None = None,
+        element_type_name: str = 'elements',
+        truncate_repr: int | None = None,
+        item_name: str | None = None,
+    ):
+        """
+        Args:
+            elements: Initial elements to add (list or dict)
+            element_type_name: Name for display (e.g., 'components', 'buses')
+            truncate_repr: Maximum number of items to show in repr. If None, show all items. Default: None
+            item_name: Singular name for error messages (e.g., 'Component', 'Carrier').
+                If None, inferred from first added item's class name.
+        """
+        super().__init__()
+        self._element_type_name = element_type_name
+        self._truncate_repr = truncate_repr
+        self._item_name = item_name
+
+        if elements is not None:
+            if isinstance(elements, dict):
+                for element in elements.values():
+                    self.add(element)
+            else:
+                for element in elements:
+                    self.add(element)
+
+    def _get_label(self, element: T) -> str:
+        """
+        Extract label from element. Must be implemented by subclasses.
+
+        Args:
+            element: Element to get label from
+
+        Returns:
+            Label string
+        """
+        raise NotImplementedError('Subclasses must implement _get_label()')
+
+    def _get_item_name(self) -> str:
+        """Get the singular item name for error messages.
+
+        Returns the explicitly set item_name, or infers from the first item's class name.
+        Falls back to 'Item' if container is empty and no name was set.
+        """
+        if self._item_name is not None:
+            return self._item_name
+        # Infer from first item's class name
+        if self:
+            first_item = next(iter(self.values()))
+            return first_item.__class__.__name__
+        return 'Item'
+
+    def add(self, element: T) -> None:
+        """Add an element to the container."""
+        label = self._get_label(element)
+        if label in self:
+            item_name = element.__class__.__name__
+            raise ValueError(
+                f'{item_name} with label "{label}" already exists in {self._element_type_name}. '
+                f'Each {item_name.lower()} must have a unique label.'
+            )
+        self[label] = element
+
+    def __setitem__(self, label: str, element: T) -> None:
+        """Set element with validation."""
+        element_label = self._get_label(element)
+        if label != element_label:
+            raise ValueError(
+                f'Key "{label}" does not match element label "{element_label}". '
+                f'Use the correct label as key or use .add() method.'
+            )
+        super().__setitem__(label, element)
+
+    def __getitem__(self, label: str) -> T:
+        """
+        Get element by label with helpful error messages.
+
+        Args:
+            label: Label of the element to retrieve
+
+        Returns:
+            The element with the given label
+
+        Raises:
+            KeyError: If element is not found, with suggestions for similar labels
+        """
+        try:
+            return super().__getitem__(label)
+        except KeyError:
+            # Provide helpful error with close matches suggestions
+            item_name = self._get_item_name()
+            suggestions = get_close_matches(label, self.keys(), n=3, cutoff=0.6)
+            error_msg = f'{item_name} "{label}" not found in {self._element_type_name}.'
+            if suggestions:
+                error_msg += f' Did you mean: {", ".join(suggestions)}?'
+            else:
+                available = list(self.keys())
+                if len(available) <= 5:
+                    error_msg += f' Available: {", ".join(available)}'
+                else:
+                    error_msg += f' Available: {", ".join(available[:5])} ... (+{len(available) - 5} more)'
+            raise KeyError(error_msg) from None
+
+    def _get_repr(self, max_items: int | None = None) -> str:
+        """
+        Get string representation with optional truncation.
+
+        Args:
+            max_items: Maximum number of items to show. If None, uses instance default (self._truncate_repr).
+                      If still None, shows all items.
+
+        Returns:
+            Formatted string representation
+        """
+        # Use provided max_items, or fall back to instance default
+        limit = max_items if max_items is not None else self._truncate_repr
+
+        count = len(self)
+        title = f'{self._element_type_name.capitalize()} ({count} item{"s" if count != 1 else ""})'
+
+        if not self:
+            r = fx_io.format_title_with_underline(title)
+            r += '<empty>\n'
+        else:
+            r = fx_io.format_title_with_underline(title)
+            sorted_names = sorted(self.keys(), key=_natural_sort_key)
+
+            if limit is not None and limit > 0 and len(sorted_names) > limit:
+                # Show truncated list
+                for name in sorted_names[:limit]:
+                    r += f' * {name}\n'
+                r += f' ... (+{len(sorted_names) - limit} more)\n'
+            else:
+                # Show all items
+                for name in sorted_names:
+                    r += f' * {name}\n'
+
+        return r
+
+    def __repr__(self) -> str:
+        """Return a string representation using the instance's truncate_repr setting."""
+        return self._get_repr()
+
+
+class ElementContainer(ContainerMixin[T]):
+    """
+    Container for Element objects (Component, Bus, Flow, Effect).
+
+    Uses element.label_full for keying.
+    """
+
+    def _get_label(self, element: T) -> str:
+        """Extract label_full from Element."""
+        return element.label_full
+
+
+class ResultsContainer(ContainerMixin[T]):
+    """
+    Container for Results objects (ComponentResults, BusResults, etc).
+
+    Uses element.label for keying.
+    """
+
+    def _get_label(self, element: T) -> str:
+        """Extract label from Results object."""
+        return element.label
+
+
+T_element = TypeVar('T_element')
+
+
+class CompositeContainerMixin(Generic[T_element]):
+    """
+    Mixin providing unified dict-like access across multiple typed containers.
+
+    This mixin enables classes that manage multiple containers (e.g., components,
+    buses, effects, flows) to provide a unified interface for accessing elements
+    across all containers, as if they were a single collection.
+
+    Type Parameter:
+        T_element: The type of elements stored in the containers. Can be a union type
+            for containers holding multiple types (e.g., 'ComponentResults | BusResults').
+
+    Key Features:
+        - Dict-like access: `obj['element_name']` searches all containers
+        - Iteration: `for label in obj:` iterates over all elements
+        - Membership: `'element' in obj` checks across all containers
+        - Standard dict methods: keys(), values(), items()
+        - Grouped display: Formatted repr showing elements by type
+        - Type hints: Full IDE and type checker support
+
+    Subclasses must implement:
+        _get_container_groups() -> dict[str, dict]:
+            Returns a dictionary mapping group names (e.g., 'Components', 'Buses')
+            to container dictionaries. Containers are displayed in the order returned.
+
+    Example:
+        ```python
+        class MySystem(CompositeContainerMixin[Component | Bus]):
+            def __init__(self):
+                self.components = {'Boiler': Component(...), 'CHP': Component(...)}
+                self.buses = {'Heat': Bus(...), 'Power': Bus(...)}
+
+            def _get_container_groups(self):
+                return {
+                    'Components': self.components,
+                    'Buses': self.buses,
+                }
+
+
+        system = MySystem()
+        comp = system['Boiler']  # Type: Component | Bus (with proper IDE support)
+        'Heat' in system  # True
+        labels = system.keys()  # Type: list[str]
+        elements = system.values()  # Type: list[Component | Bus]
+        ```
+
+    Integration with ContainerMixin:
+        This mixin is designed to work alongside ContainerMixin-based containers
+        (ElementContainer, ResultsContainer) by aggregating them into a unified
+        interface while preserving their individual functionality.
+    """
+
+    def _get_container_groups(self) -> dict[str, ContainerMixin[Any]]:
+        """
+        Return ordered dict of container groups to aggregate.
+
+        Returns:
+            Dictionary mapping group names to container objects (e.g., ElementContainer, ResultsContainer).
+            Group names should be capitalized (e.g., 'Components', 'Buses').
+            Order determines display order in __repr__.
+
+        Example:
+            ```python
+            return {
+                'Components': self.components,
+                'Buses': self.buses,
+                'Effects': self.effects,
+            }
+            ```
+        """
+        raise NotImplementedError('Subclasses must implement _get_container_groups()')
+
+    def __getitem__(self, key: str) -> T_element:
+        """
+        Get element by label, searching all containers.
+
+        Args:
+            key: Element label to find
+
+        Returns:
+            The element with the given label
+
+        Raises:
+            KeyError: If element not found, with helpful suggestions
+        """
+        # Search all containers in order
+        for container in self._get_container_groups().values():
+            if key in container:
+                return container[key]
+
+        # Element not found - provide helpful error
+        all_elements = {}
+        for container in self._get_container_groups().values():
+            all_elements.update(container)
+
+        suggestions = get_close_matches(key, all_elements.keys(), n=3, cutoff=0.6)
+        error_msg = f'Element "{key}" not found.'
+
+        if suggestions:
+            error_msg += f' Did you mean: {", ".join(suggestions)}?'
+        else:
+            available = list(all_elements.keys())
+            if len(available) <= 5:
+                error_msg += f' Available: {", ".join(available)}'
+            else:
+                error_msg += f' Available: {", ".join(available[:5])} ... (+{len(available) - 5} more)'
+
+        raise KeyError(error_msg)
+
+    def __iter__(self):
+        """Iterate over all element labels across all containers."""
+        for container in self._get_container_groups().values():
+            yield from container.keys()
+
+    def __len__(self) -> int:
+        """Return total count of elements across all containers."""
+        return sum(len(container) for container in self._get_container_groups().values())
+
+    def __contains__(self, key: str) -> bool:
+        """Check if element exists in any container."""
+        return any(key in container for container in self._get_container_groups().values())
+
+    def keys(self) -> list[str]:
+        """Return all element labels across all containers."""
+        return list(self)
+
+    def values(self) -> list[T_element]:
+        """Return all element objects across all containers."""
+        vals = []
+        for container in self._get_container_groups().values():
+            vals.extend(container.values())
+        return vals
+
+    def items(self) -> list[tuple[str, T_element]]:
+        """Return (label, element) pairs for all elements."""
+        items = []
+        for container in self._get_container_groups().values():
+            items.extend(container.items())
+        return items
+
+    def _format_grouped_containers(self, title: str | None = None) -> str:
+        """
+        Format containers as grouped string representation using each container's repr.
+
+        Args:
+            title: Optional title for the representation. If None, no title is shown.
+
+        Returns:
+            Formatted string with groups and their elements.
+            Empty groups are automatically hidden.
+
+        Example output:
+            ```
+            Components (1 item)
+            -------------------
+             * Boiler
+
+            Buses (2 items)
+            ---------------
+             * Heat
+             * Power
+            ```
+        """
+        parts = []
+
+        if title:
+            parts.append(fx_io.format_title_with_underline(title))
+
+        container_groups = self._get_container_groups()
+        for container in container_groups.values():
+            if container:  # Only show non-empty groups
+                if parts:  # Add spacing between sections
+                    parts.append('')
+                # Use container's __repr__ which respects its truncate_repr setting
+                parts.append(repr(container).rstrip('\n'))
+
+        return '\n'.join(parts)
+
+
 class Submodel(SubmodelsMixin):
     """Stores Variables and Constraints. Its a subset of a FlowSystemModel.
     Variables and constraints are stored in the main FlowSystemModel, and are referenced here.
@@ -924,8 +1739,22 @@ class Submodel(SubmodelsMixin):
         logger.debug(f'Creating {self.__class__.__name__}  "{self.label_full}"')
         self._do_modeling()
 
-    def add_variables(self, short_name: str = None, **kwargs) -> linopy.Variable:
-        """Create and register a variable in one step"""
+    def add_variables(
+        self,
+        short_name: str = None,
+        category: VariableCategory = None,
+        **kwargs: Any,
+    ) -> linopy.Variable:
+        """Create and register a variable in one step.
+
+        Args:
+            short_name: Short name for the variable (used as suffix in full name).
+            category: Category for segment expansion handling. See VariableCategory.
+            **kwargs: Additional arguments passed to linopy.Model.add_variables().
+
+        Returns:
+            The created linopy Variable.
+        """
         if kwargs.get('name') is None:
             if short_name is None:
                 raise ValueError('Short name must be provided when no name is given')
@@ -933,6 +1762,11 @@ class Submodel(SubmodelsMixin):
 
         variable = self._model.add_variables(**kwargs)
         self.register_variable(variable, short_name)
+
+        # Register category in FlowSystemModel for segment expansion handling
+        if category is not None:
+            self._model.variable_categories[variable.name] = category
+
         return variable
 
     def add_constraints(self, expression, short_name: str = None, **kwargs) -> linopy.Constraint:
@@ -1061,9 +1895,7 @@ class Submodel(SubmodelsMixin):
         }
 
         # Format sections with headers and underlines
-        formatted_sections = []
-        for section_header, section_content in sections.items():
-            formatted_sections.append(f'{section_header}\n{"-" * len(section_header)}\n{section_content}')
+        formatted_sections = fx_io.format_sections_with_headers(sections)
 
         model_string = f'Submodel "{self.label_of_model}":'
         all_sections = '\n'.join(formatted_sections)
@@ -1071,11 +1903,16 @@ class Submodel(SubmodelsMixin):
         return f'{model_string}\n{"=" * len(model_string)}\n\n{all_sections}'
 
     @property
-    def hours_per_step(self):
-        return self._model.hours_per_step
+    def timestep_duration(self):
+        return self._model.timestep_duration
 
     def _do_modeling(self):
-        """Called at the end of initialization. Override in subclasses to create variables and constraints."""
+        """
+        Override in subclasses to create variables, constraints, and submodels.
+
+        This method is called during __init__. Create all nested submodels first
+        (so their variables exist), then create constraints that reference those variables.
+        """
         pass
 
 
@@ -1107,7 +1944,7 @@ class Submodels:
     def __repr__(self) -> str:
         """Simple representation of the submodels collection."""
         if not self.data:
-            return 'flixopt.structure.Submodels:\n----------------------------\n <empty>\n'
+            return fx_io.format_title_with_underline('flixopt.structure.Submodels') + ' <empty>\n'
 
         total_vars = sum(len(submodel.variables) for submodel in self.data.values())
         total_cons = sum(len(submodel.constraints) for submodel in self.data.values())
@@ -1115,18 +1952,15 @@ class Submodels:
         title = (
             f'flixopt.structure.Submodels ({total_vars} vars, {total_cons} constraints, {len(self.data)} submodels):'
         )
-        underline = '-' * len(title)
 
-        if not self.data:
-            return f'{title}\n{underline}\n <empty>\n'
-        sub_models_string = ''
+        result = fx_io.format_title_with_underline(title)
         for name, submodel in self.data.items():
             type_name = submodel.__class__.__name__
             var_count = len(submodel.variables)
             con_count = len(submodel.constraints)
-            sub_models_string += f'\n * {name} [{type_name}] ({var_count}v/{con_count}c)'
+            result += f' * {name} [{type_name}] ({var_count}v/{con_count}c)\n'
 
-        return f'{title}\n{underline}{sub_models_string}\n'
+        return result
 
     def items(self) -> ItemsView[str, Submodel]:
         return self.data.items()