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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of flixopt might be problematic. Click here for more details.

flixopt/results.py CHANGED
@@ -2,7 +2,8 @@ import datetime
2
2
  import json
3
3
  import logging
4
4
  import pathlib
5
- from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union
5
+ import warnings
6
+ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union
6
7
 
7
8
  import linopy
8
9
  import matplotlib.pyplot as plt
@@ -14,7 +15,8 @@ import yaml
14
15
 
15
16
  from . import io as fx_io
16
17
  from . import plotting
17
- from .core import TimeSeriesCollection
18
+ from .core import DataConverter, TimeSeriesCollection
19
+ from .flow_system import FlowSystem
18
20
 
19
21
  if TYPE_CHECKING:
20
22
  import pyvis
@@ -25,6 +27,11 @@ if TYPE_CHECKING:
25
27
  logger = logging.getLogger('flixopt')
26
28
 
27
29
 
30
+ class _FlowSystemRestorationError(Exception):
31
+ """Exception raised when a FlowSystem cannot be restored from dataset."""
32
+ pass
33
+
34
+
28
35
  class CalculationResults:
29
36
  """Results container for Calculation results.
30
37
 
@@ -37,7 +44,7 @@ class CalculationResults:
37
44
 
38
45
  Attributes:
39
46
  solution (xr.Dataset): Dataset containing optimization results.
40
- flow_system (xr.Dataset): Dataset containing the flow system.
47
+ flow_system_data (xr.Dataset): Dataset containing the flow system.
41
48
  summary (Dict): Information about the calculation.
42
49
  name (str): Name identifier for the calculation.
43
50
  model (linopy.Model): The optimization model (if available).
@@ -92,7 +99,7 @@ class CalculationResults:
92
99
 
93
100
  return cls(
94
101
  solution=fx_io.load_dataset_from_netcdf(paths.solution),
95
- flow_system=fx_io.load_dataset_from_netcdf(paths.flow_system),
102
+ flow_system_data=fx_io.load_dataset_from_netcdf(paths.flow_system),
96
103
  name=name,
97
104
  folder=folder,
98
105
  model=model,
@@ -118,7 +125,7 @@ class CalculationResults:
118
125
  """
119
126
  return cls(
120
127
  solution=calculation.model.solution,
121
- flow_system=calculation.flow_system.as_dataset(constants_in_dataset=True),
128
+ flow_system_data=calculation.flow_system.as_dataset(constants_in_dataset=True),
122
129
  summary=calculation.summary,
123
130
  model=calculation.model,
124
131
  name=calculation.name,
@@ -128,47 +135,84 @@ class CalculationResults:
128
135
  def __init__(
129
136
  self,
130
137
  solution: xr.Dataset,
131
- flow_system: xr.Dataset,
138
+ flow_system_data: xr.Dataset,
132
139
  name: str,
133
140
  summary: Dict,
134
141
  folder: Optional[pathlib.Path] = None,
135
142
  model: Optional[linopy.Model] = None,
143
+ **kwargs, # To accept old "flow_system" parameter
136
144
  ):
137
145
  """
138
146
  Args:
139
147
  solution: The solution of the optimization.
140
- flow_system: The flow_system that was used to create the calculation as a datatset.
148
+ flow_system_data: The flow_system that was used to create the calculation as a datatset.
141
149
  name: The name of the calculation.
142
150
  summary: Information about the calculation,
143
151
  folder: The folder where the results are saved.
144
152
  model: The linopy model that was used to solve the calculation.
153
+ Deprecated:
154
+ flow_system: Use flow_system_data instead.
145
155
  """
156
+ # Handle potential old "flow_system" parameter for backward compatibility
157
+ if 'flow_system' in kwargs and flow_system_data is None:
158
+ flow_system_data = kwargs.pop('flow_system')
159
+ warnings.warn(
160
+ "The 'flow_system' parameter is deprecated. Use 'flow_system_data' instead."
161
+ "Acess is now by '.flow_system_data', while '.flow_system' returns the restored FlowSystem.",
162
+ DeprecationWarning,
163
+ stacklevel=2,
164
+ )
165
+
146
166
  self.solution = solution
147
- self.flow_system = flow_system
167
+ self.flow_system_data = flow_system_data
148
168
  self.summary = summary
149
169
  self.name = name
150
170
  self.model = model
151
171
  self.folder = pathlib.Path(folder) if folder is not None else pathlib.Path.cwd() / 'results'
152
172
  self.components = {
153
- label: ComponentResults.from_json(self, infos) for label, infos in self.solution.attrs['Components'].items()
173
+ label: ComponentResults(self, **infos) for label, infos in self.solution.attrs['Components'].items()
154
174
  }
155
175
 
156
- self.buses = {label: BusResults.from_json(self, infos) for label, infos in self.solution.attrs['Buses'].items()}
176
+ self.buses = {label: BusResults(self, **infos) for label, infos in self.solution.attrs['Buses'].items()}
157
177
 
158
178
  self.effects = {
159
- label: EffectResults.from_json(self, infos) for label, infos in self.solution.attrs['Effects'].items()
179
+ label: EffectResults(self, **infos) for label, infos in self.solution.attrs['Effects'].items()
160
180
  }
161
181
 
182
+ if 'Flows' not in self.solution.attrs:
183
+ warnings.warn(
184
+ 'No Data about flows found in the results. This data is only included since v2.2.0. Some functionality '
185
+ 'is not availlable. We recommend to evaluate your results with a version <2.2.0.',
186
+ stacklevel=2,
187
+ )
188
+ self.flows = {}
189
+ else:
190
+ self.flows = {
191
+ label: FlowResults(self, **infos) for label, infos in self.solution.attrs.get('Flows', {}).items()
192
+ }
193
+
162
194
  self.timesteps_extra = self.solution.indexes['time']
163
195
  self.hours_per_timestep = TimeSeriesCollection.calculate_hours_per_timestep(self.timesteps_extra)
196
+ self.scenarios = self.solution.indexes['scenario'] if 'scenario' in self.solution.indexes else None
197
+
198
+ self._effect_share_factors = None
199
+ self._flow_system = None
200
+
201
+ self._flow_rates = None
202
+ self._flow_hours = None
203
+ self._sizes = None
204
+ self._effects_per_component = {'operation': None, 'invest': None, 'total': None}
205
+ self._flow_network_info_ = None
164
206
 
165
- def __getitem__(self, key: str) -> Union['ComponentResults', 'BusResults', 'EffectResults']:
207
+ def __getitem__(self, key: str) -> Union['ComponentResults', 'BusResults', 'EffectResults', 'FlowResults']:
166
208
  if key in self.components:
167
209
  return self.components[key]
168
210
  if key in self.buses:
169
211
  return self.buses[key]
170
212
  if key in self.effects:
171
213
  return self.effects[key]
214
+ if key in self.flows:
215
+ return self.flows[key]
172
216
  raise KeyError(f'No element with label {key} found.')
173
217
 
174
218
  @property
@@ -195,20 +239,376 @@ class CalculationResults:
195
239
  raise ValueError('The linopy model is not available.')
196
240
  return self.model.constraints
197
241
 
242
+ @property
243
+ def effect_share_factors(self):
244
+ if self._effect_share_factors is None:
245
+ effect_share_factors = self.flow_system.effects.calculate_effect_share_factors()
246
+ self._effect_share_factors = {'operation': effect_share_factors[0],
247
+ 'invest': effect_share_factors[1]}
248
+ return self._effect_share_factors
249
+
250
+ @property
251
+ def flow_system(self) -> 'FlowSystem':
252
+ """ The restored flow_system that was used to create the calculation.
253
+ Contains all input parameters."""
254
+ if self._flow_system is None:
255
+ try:
256
+ from . import FlowSystem
257
+ current_logger_level = logger.getEffectiveLevel()
258
+ logger.setLevel(logging.CRITICAL)
259
+ self._flow_system = FlowSystem.from_dataset(self.flow_system_data)
260
+ self._flow_system._connect_network()
261
+ logger.setLevel(current_logger_level)
262
+ except Exception as e:
263
+ logger.critical(f'Not able to restore FlowSystem from dataset. Some functionality is not availlable. {e}')
264
+ raise _FlowSystemRestorationError(f'Not able to restore FlowSystem from dataset. {e}') from e
265
+ return self._flow_system
266
+
198
267
  def filter_solution(
199
- self, variable_dims: Optional[Literal['scalar', 'time']] = None, element: Optional[str] = None
268
+ self,
269
+ variable_dims: Optional[Literal['scalar', 'time', 'scenario', 'timeonly', 'scenarioonly']] = None,
270
+ element: Optional[str] = None,
271
+ timesteps: Optional[pd.DatetimeIndex] = None,
272
+ scenarios: Optional[pd.Index] = None,
273
+ contains: Optional[Union[str, List[str]]] = None,
274
+ startswith: Optional[Union[str, List[str]]] = None,
200
275
  ) -> xr.Dataset:
201
276
  """
202
277
  Filter the solution to a specific variable dimension and element.
203
278
  If no element is specified, all elements are included.
204
279
 
205
280
  Args:
206
- variable_dims: The dimension of the variables to filter for.
281
+ variable_dims: The dimension of which to get variables from.
282
+ - 'scalar': Get scalar variables (without dimensions)
283
+ - 'time': Get time-dependent variables (with a time dimension)
284
+ - 'scenario': Get scenario-dependent variables (with ONLY a scenario dimension)
285
+ - 'timeonly': Get time-dependent variables (with ONLY a time dimension)
286
+ - 'scenarioonly': Get scenario-dependent variables (with ONLY a scenario dimension)
207
287
  element: The element to filter for.
288
+ timesteps: Optional time indexes to select. Can be:
289
+ - pd.DatetimeIndex: Multiple timesteps
290
+ - str/pd.Timestamp: Single timestep
291
+ Defaults to all available timesteps.
292
+ scenarios: Optional scenario indexes to select. Can be:
293
+ - pd.Index: Multiple scenarios
294
+ - str/int: Single scenario (int is treated as a label, not an index position)
295
+ Defaults to all available scenarios.
296
+ contains: Filter variables that contain this string or strings.
297
+ If a list is provided, variables must contain ALL strings in the list.
298
+ startswith: Filter variables that start with this string or strings.
299
+ If a list is provided, variables must start with ANY of the strings in the list.
300
+ """
301
+ return filter_dataset(
302
+ self.solution if element is None else self[element].solution,
303
+ variable_dims=variable_dims,
304
+ timesteps=timesteps,
305
+ scenarios=scenarios,
306
+ contains=contains,
307
+ startswith=startswith,
308
+ )
309
+
310
+ def effects_per_component(self, mode: Literal['operation', 'invest', 'total'] = 'total') -> xr.Dataset:
311
+ """Returns a dataset containing effect totals for each components (including their flows).
312
+
313
+ Args:
314
+ mode: Which effects to contain. (operation, invest, total)
315
+
316
+ Returns:
317
+ An xarray Dataset with an additional component dimension and effects as variables.
318
+ """
319
+ if mode not in ['operation', 'invest', 'total']:
320
+ raise ValueError(f'Invalid mode {mode}')
321
+ if self._effects_per_component[mode] is None:
322
+ self._effects_per_component[mode] = self._create_effects_dataset(mode)
323
+ return self._effects_per_component[mode]
324
+
325
+ def flow_rates(
326
+ self,
327
+ start: Optional[Union[str, List[str]]] = None,
328
+ end: Optional[Union[str, List[str]]] = None,
329
+ component: Optional[Union[str, List[str]]] = None,
330
+ ) -> xr.DataArray:
331
+ """Returns a DataArray containing the flow rates of each Flow.
332
+
333
+ Args:
334
+ start: Optional source node(s) to filter by. Can be a single node name or a list of names.
335
+ end: Optional destination node(s) to filter by. Can be a single node name or a list of names.
336
+ component: Optional component(s) to filter by. Can be a single component name or a list of names.
337
+
338
+ Further usage:
339
+ Convert the dataarray to a dataframe:
340
+ >>>results.flow_rates().to_pandas()
341
+ Get the max or min over time:
342
+ >>>results.flow_rates().max('time')
343
+ Sum up the flow rates of flows with the same start and end:
344
+ >>>results.flow_rates(end='Fernwärme').groupby('start').sum(dim='flow')
345
+ To recombine filtered dataarrays, use `xr.concat` with dim 'flow':
346
+ >>>xr.concat([results.flow_rates(start='Fernwärme'), results.flow_rates(end='Fernwärme')], dim='flow')
347
+ """
348
+ if self._flow_rates is None:
349
+ self._flow_rates = self._assign_flow_coords(
350
+ xr.concat([flow.flow_rate.rename(flow.label) for flow in self.flows.values()],
351
+ dim=pd.Index(self.flows.keys(), name='flow'))
352
+ ).rename('flow_rates')
353
+ filters = {k: v for k, v in {'start': start, 'end': end, 'component': component}.items() if v is not None}
354
+ return filter_dataarray_by_coord(self._flow_rates, **filters)
355
+
356
+ def flow_hours(
357
+ self,
358
+ start: Optional[Union[str, List[str]]] = None,
359
+ end: Optional[Union[str, List[str]]] = None,
360
+ component: Optional[Union[str, List[str]]] = None,
361
+ ) -> xr.DataArray:
362
+ """Returns a DataArray containing the flow hours of each Flow.
363
+
364
+ Flow hours represent the total energy/material transferred over time,
365
+ calculated by multiplying flow rates by the duration of each timestep.
366
+
367
+ Args:
368
+ start: Optional source node(s) to filter by. Can be a single node name or a list of names.
369
+ end: Optional destination node(s) to filter by. Can be a single node name or a list of names.
370
+ component: Optional component(s) to filter by. Can be a single component name or a list of names.
371
+
372
+ Further usage:
373
+ Convert the dataarray to a dataframe:
374
+ >>>results.flow_hours().to_pandas()
375
+ Sum up the flow hours over time:
376
+ >>>results.flow_hours().sum('time')
377
+ Sum up the flow hours of flows with the same start and end:
378
+ >>>results.flow_hours(end='Fernwärme').groupby('start').sum(dim='flow')
379
+ To recombine filtered dataarrays, use `xr.concat` with dim 'flow':
380
+ >>>xr.concat([results.flow_hours(start='Fernwärme'), results.flow_hours(end='Fernwärme')], dim='flow')
381
+
382
+ """
383
+ if self._flow_hours is None:
384
+ self._flow_hours = (self.flow_rates() * self.hours_per_timestep).rename('flow_hours')
385
+ filters = {k: v for k, v in {'start': start, 'end': end, 'component': component}.items() if v is not None}
386
+ return filter_dataarray_by_coord(self._flow_hours, **filters)
387
+
388
+ def sizes(
389
+ self,
390
+ start: Optional[Union[str, List[str]]] = None,
391
+ end: Optional[Union[str, List[str]]] = None,
392
+ component: Optional[Union[str, List[str]]] = None
393
+ ) -> xr.DataArray:
394
+ """Returns a dataset with the sizes of the Flows.
395
+ Args:
396
+ start: Optional source node(s) to filter by. Can be a single node name or a list of names.
397
+ end: Optional destination node(s) to filter by. Can be a single node name or a list of names.
398
+ component: Optional component(s) to filter by. Can be a single component name or a list of names.
399
+
400
+ Further usage:
401
+ Convert the dataarray to a dataframe:
402
+ >>>results.sizes().to_pandas()
403
+ To recombine filtered dataarrays, use `xr.concat` with dim 'flow':
404
+ >>>xr.concat([results.sizes(start='Fernwärme'), results.sizes(end='Fernwärme')], dim='flow')
405
+
406
+ """
407
+ if self._sizes is None:
408
+ self._sizes = self._assign_flow_coords(
409
+ xr.concat([flow.size.rename(flow.label) for flow in self.flows.values()],
410
+ dim=pd.Index(self.flows.keys(), name='flow'))
411
+ ).rename('flow_sizes')
412
+ filters = {k: v for k, v in {'start': start, 'end': end, 'component': component}.items() if v is not None}
413
+ return filter_dataarray_by_coord(self._sizes, **filters)
414
+
415
+ def _assign_flow_coords(self, da: xr.DataArray):
416
+ # Add start and end coordinates
417
+ da = da.assign_coords({
418
+ 'start': ('flow', [flow.start for flow in self.flows.values()]),
419
+ 'end': ('flow', [flow.end for flow in self.flows.values()]),
420
+ 'component': ('flow', [flow.component for flow in self.flows.values()]),
421
+ })
422
+
423
+ # Ensure flow is the last dimension if needed
424
+ existing_dims = [d for d in da.dims if d != 'flow']
425
+ da = da.transpose(*(existing_dims + ['flow']))
426
+ return da
427
+
428
+ def _get_flow_network_info(self) -> Dict[str, Dict[str, str]]:
429
+ flow_network_info = {}
430
+
431
+ for flow in self.flows.values():
432
+ flow_network_info[flow.label] = {
433
+ 'label': flow.label,
434
+ 'start': flow.start,
435
+ 'end': flow.end,
436
+ }
437
+ return flow_network_info
438
+
439
+ def get_effect_shares(
440
+ self,
441
+ element: str,
442
+ effect: str,
443
+ mode: Optional[Literal['operation', 'invest']] = None,
444
+ include_flows: bool = False
445
+ ) -> xr.Dataset:
446
+ """Retrieves individual effect shares for a specific element and effect.
447
+ Either for operation, investment, or both modes combined.
448
+ Only includes the direct shares.
449
+
450
+ Args:
451
+ element: The element identifier for which to retrieve effect shares.
452
+ effect: The effect identifier for which to retrieve shares.
453
+ mode: Optional. The mode to retrieve shares for. Can be 'operation', 'invest',
454
+ or None to retrieve both. Defaults to None.
455
+
456
+ Returns:
457
+ An xarray Dataset containing the requested effect shares. If mode is None,
458
+ returns a merged Dataset containing both operation and investment shares.
459
+
460
+ Raises:
461
+ ValueError: If the specified effect is not available or if mode is invalid.
462
+ """
463
+ if effect not in self.effects:
464
+ raise ValueError(f'Effect {effect} is not available.')
465
+
466
+ if mode is None:
467
+ return xr.merge([self.get_effect_shares(element=element, effect=effect, mode='operation', include_flows=include_flows),
468
+ self.get_effect_shares(element=element, effect=effect, mode='invest', include_flows=include_flows)])
469
+
470
+ if mode not in ['operation', 'invest']:
471
+ raise ValueError(f'Mode {mode} is not available. Choose between "operation" and "invest".')
472
+
473
+ ds = xr.Dataset()
474
+
475
+ label = f'{element}->{effect}({mode})'
476
+ if label in self.solution:
477
+ ds = xr.Dataset({label: self.solution[label]})
478
+
479
+ if include_flows:
480
+ if element not in self.components:
481
+ raise ValueError(f'Only use Components when retrieving Effects including flows. Got {element}')
482
+ flows = [label.split('|')[0] for label in self.components[element].inputs + self.components[element].outputs]
483
+ return xr.merge(
484
+ [ds] + [self.get_effect_shares(element=flow, effect=effect, mode=mode, include_flows=False)
485
+ for flow in flows]
486
+ )
487
+
488
+ return ds
489
+
490
+ def _compute_effect_total(
491
+ self,
492
+ element: str,
493
+ effect: str,
494
+ mode: Literal['operation', 'invest', 'total'] = 'total',
495
+ include_flows: bool = False,
496
+ ) -> xr.DataArray:
497
+ """Calculates the total effect for a specific element and effect.
498
+
499
+ This method computes the total direct and indirect effects for a given element
500
+ and effect, considering the conversion factors between different effects.
501
+
502
+ Args:
503
+ element: The element identifier for which to calculate total effects.
504
+ effect: The effect identifier to calculate.
505
+ mode: The calculation mode. Options are:
506
+ 'operation': Returns operation-specific effects.
507
+ 'invest': Returns investment-specific effects.
508
+ 'total': Returns the sum of operation effects (across all timesteps)
509
+ and investment effects. Defaults to 'total'.
510
+ include_flows: Whether to include effects from flows connected to this element.
511
+
512
+ Returns:
513
+ An xarray DataArray containing the total effects, named with pattern
514
+ '{element}->{effect}' for mode='total' or '{element}->{effect}({mode})'
515
+ for other modes.
516
+
517
+ Raises:
518
+ ValueError: If the specified effect is not available.
208
519
  """
209
- if element is not None:
210
- return filter_dataset(self[element].solution, variable_dims)
211
- return filter_dataset(self.solution, variable_dims)
520
+ if effect not in self.effects:
521
+ raise ValueError(f'Effect {effect} is not available.')
522
+
523
+ if mode == 'total':
524
+ operation = self._compute_effect_total(element=element, effect=effect, mode='operation', include_flows=include_flows)
525
+ invest = self._compute_effect_total(element=element, effect=effect, mode='invest', include_flows=include_flows)
526
+ if invest.isnull().all() and operation.isnull().all():
527
+ return xr.DataArray(np.nan)
528
+ if operation.isnull().all():
529
+ return invest.rename(f'{element}->{effect}')
530
+ operation = operation.sum('time')
531
+ if invest.isnull().all():
532
+ return operation.rename(f'{element}->{effect}')
533
+ if 'time' in operation.indexes:
534
+ operation = operation.sum('time')
535
+ return invest + operation
536
+
537
+ total = xr.DataArray(0)
538
+ share_exists = False
539
+
540
+ relevant_conversion_factors = {
541
+ key[0]: value for key, value in self.effect_share_factors[mode].items() if key[1] == effect
542
+ }
543
+ relevant_conversion_factors[effect] = 1 # Share to itself is 1
544
+
545
+ for target_effect, conversion_factor in relevant_conversion_factors.items():
546
+ label = f'{element}->{target_effect}({mode})'
547
+ if label in self.solution:
548
+ share_exists = True
549
+ da = self.solution[label]
550
+ total = da * conversion_factor + total
551
+
552
+ if include_flows:
553
+ if element not in self.components:
554
+ raise ValueError(f'Only use Components when retrieving Effects including flows. Got {element}')
555
+ flows = [label.split('|')[0] for label in
556
+ self.components[element].inputs + self.components[element].outputs]
557
+ for flow in flows:
558
+ label = f'{flow}->{target_effect}({mode})'
559
+ if label in self.solution:
560
+ share_exists = True
561
+ da = self.solution[label]
562
+ total = da * conversion_factor + total
563
+ if not share_exists:
564
+ total = xr.DataArray(np.nan)
565
+ return total.rename(f'{element}->{effect}({mode})')
566
+
567
+ def _create_effects_dataset(self, mode: Literal['operation', 'invest', 'total'] = 'total') -> xr.Dataset:
568
+ """Creates a dataset containing effect totals for all components (including their flows).
569
+ The dataset does contain the direct as well as the indirect effects of each component.
570
+
571
+ Args:
572
+ mode: The calculation mode ('operation', 'invest', or 'total').
573
+
574
+ Returns:
575
+ An xarray Dataset with components as dimension and effects as variables.
576
+ """
577
+ # Create an empty dataset
578
+ ds = xr.Dataset()
579
+
580
+ # Add each effect as a variable to the dataset
581
+ for effect in self.effects:
582
+ # Create a list of DataArrays, one for each component
583
+ component_arrays = [
584
+ self._compute_effect_total(element=component, effect=effect, mode=mode, include_flows=True).expand_dims(
585
+ component=[component]
586
+ ) # Add component dimension to each array
587
+ for component in list(self.components)
588
+ ]
589
+
590
+ # Combine all components into one DataArray for this effect
591
+ if component_arrays:
592
+ effect_array = xr.concat(component_arrays, dim='component', coords='minimal')
593
+ # Add this effect as a variable to the dataset
594
+ ds[effect] = effect_array
595
+
596
+ # For now include a test to ensure correctness
597
+ suffix = {
598
+ 'operation': '(operation)|total_per_timestep',
599
+ 'invest': '(invest)|total',
600
+ 'total': '|total',
601
+ }
602
+ for effect in self.effects:
603
+ label = f'{effect}{suffix[mode]}'
604
+ computed = ds[effect].sum('component')
605
+ found = self.solution[label]
606
+ if not np.allclose(computed.values, found.fillna(0).values):
607
+ logger.critical(
608
+ f'Results for {effect}({mode}) in effects_dataset doesnt match {label}\n{computed=}\n, {found=}'
609
+ )
610
+
611
+ return ds
212
612
 
213
613
  def plot_heatmap(
214
614
  self,
@@ -219,10 +619,32 @@ class CalculationResults:
219
619
  save: Union[bool, pathlib.Path] = False,
220
620
  show: bool = True,
221
621
  engine: plotting.PlottingEngine = 'plotly',
622
+ scenario: Optional[Union[str, int]] = None,
222
623
  ) -> Union[plotly.graph_objs.Figure, Tuple[plt.Figure, plt.Axes]]:
624
+ """
625
+ Plots a heatmap of the solution of a variable.
626
+
627
+ Args:
628
+ variable_name: The name of the variable to plot.
629
+ heatmap_timeframes: The timeframes to use for the heatmap.
630
+ heatmap_timesteps_per_frame: The timesteps per frame to use for the heatmap.
631
+ color_map: The color map to use for the heatmap.
632
+ save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
633
+ show: Whether to show the plot or not.
634
+ engine: The engine to use for plotting. Can be either 'plotly' or 'matplotlib'.
635
+ scenario: The scenario to plot. Defaults to the first scenario. Has no effect without scenarios present
636
+ """
637
+ dataarray = self.solution[variable_name]
638
+
639
+ scenario_suffix = ''
640
+ if 'scenario' in dataarray.indexes:
641
+ chosen_scenario = scenario or self.scenarios[0]
642
+ dataarray = dataarray.sel(scenario=chosen_scenario).drop_vars('scenario')
643
+ scenario_suffix = f'--{chosen_scenario}'
644
+
223
645
  return plot_heatmap(
224
- dataarray=self.solution[variable_name],
225
- name=variable_name,
646
+ dataarray=dataarray,
647
+ name=f'{variable_name}{scenario_suffix}',
226
648
  folder=self.folder,
227
649
  heatmap_timeframes=heatmap_timeframes,
228
650
  heatmap_timesteps_per_frame=heatmap_timesteps_per_frame,
@@ -244,16 +666,9 @@ class CalculationResults:
244
666
  show: bool = False,
245
667
  ) -> 'pyvis.network.Network':
246
668
  """See flixopt.flow_system.FlowSystem.plot_network"""
247
- try:
248
- from .flow_system import FlowSystem
249
-
250
- flow_system = FlowSystem.from_dataset(self.flow_system)
251
- except Exception as e:
252
- logger.critical(f'Could not reconstruct the flow_system from dataset: {e}')
253
- return None
254
669
  if path is None:
255
670
  path = self.folder / f'{self.name}--network.html'
256
- return flow_system.plot_network(controls=controls, path=path, show=show)
671
+ return self.flow_system.plot_network(controls=controls, path=path, show=show)
257
672
 
258
673
  def to_file(
259
674
  self,
@@ -286,7 +701,7 @@ class CalculationResults:
286
701
  paths = fx_io.CalculationResultsPaths(folder, name)
287
702
 
288
703
  fx_io.save_dataset_to_netcdf(self.solution, paths.solution, compression=compression)
289
- fx_io.save_dataset_to_netcdf(self.flow_system, paths.flow_system, compression=compression)
704
+ fx_io.save_dataset_to_netcdf(self.flow_system_data, paths.flow_system, compression=compression)
290
705
 
291
706
  with open(paths.summary, 'w', encoding='utf-8') as f:
292
707
  yaml.dump(self.summary, f, allow_unicode=True, sort_keys=False, indent=4, width=1000)
@@ -307,10 +722,6 @@ class CalculationResults:
307
722
 
308
723
 
309
724
  class _ElementResults:
310
- @classmethod
311
- def from_json(cls, calculation_results, json_data: Dict) -> '_ElementResults':
312
- return cls(calculation_results, json_data['label'], json_data['variables'], json_data['constraints'])
313
-
314
725
  def __init__(
315
726
  self, calculation_results: CalculationResults, label: str, variables: List[str], constraints: List[str]
316
727
  ):
@@ -343,30 +754,51 @@ class _ElementResults:
343
754
  """
344
755
  if self._calculation_results.model is None:
345
756
  raise ValueError('The linopy model is not available.')
346
- return self._calculation_results.model.constraints[self._variable_names]
757
+ return self._calculation_results.model.constraints[self._constraint_names]
347
758
 
348
- def filter_solution(self, variable_dims: Optional[Literal['scalar', 'time']] = None) -> xr.Dataset:
759
+ def filter_solution(
760
+ self,
761
+ variable_dims: Optional[Literal['scalar', 'time', 'scenario', 'timeonly', 'scenarioonly']] = None,
762
+ timesteps: Optional[pd.DatetimeIndex] = None,
763
+ scenarios: Optional[pd.Index] = None,
764
+ contains: Optional[Union[str, List[str]]] = None,
765
+ startswith: Optional[Union[str, List[str]]] = None,
766
+ ) -> xr.Dataset:
349
767
  """
350
- Filter the solution of the element by dimension.
768
+ Filter the solution to a specific variable dimension and element.
769
+ If no element is specified, all elements are included.
351
770
 
352
771
  Args:
353
- variable_dims: The dimension of the variables to filter for.
772
+ variable_dims: The dimension of which to get variables from.
773
+ - 'scalar': Get scalar variables (without dimensions)
774
+ - 'time': Get time-dependent variables (with a time dimension)
775
+ - 'scenario': Get scenario-dependent variables (with ONLY a scenario dimension)
776
+ - 'timeonly': Get time-dependent variables (with ONLY a time dimension)
777
+ - 'scenarioonly': Get scenario-dependent variables (with ONLY a scenario dimension)
778
+ timesteps: Optional time indexes to select. Can be:
779
+ - pd.DatetimeIndex: Multiple timesteps
780
+ - str/pd.Timestamp: Single timestep
781
+ Defaults to all available timesteps.
782
+ scenarios: Optional scenario indexes to select. Can be:
783
+ - pd.Index: Multiple scenarios
784
+ - str/int: Single scenario (int is treated as a label, not an index position)
785
+ Defaults to all available scenarios.
786
+ contains: Filter variables that contain this string or strings.
787
+ If a list is provided, variables must contain ALL strings in the list.
788
+ startswith: Filter variables that start with this string or strings.
789
+ If a list is provided, variables must start with ANY of the strings in the list.
354
790
  """
355
- return filter_dataset(self.solution, variable_dims)
791
+ return filter_dataset(
792
+ self.solution,
793
+ variable_dims=variable_dims,
794
+ timesteps=timesteps,
795
+ scenarios=scenarios,
796
+ contains=contains,
797
+ startswith=startswith,
798
+ )
356
799
 
357
800
 
358
801
  class _NodeResults(_ElementResults):
359
- @classmethod
360
- def from_json(cls, calculation_results, json_data: Dict) -> '_NodeResults':
361
- return cls(
362
- calculation_results,
363
- json_data['label'],
364
- json_data['variables'],
365
- json_data['constraints'],
366
- json_data['inputs'],
367
- json_data['outputs'],
368
- )
369
-
370
802
  def __init__(
371
803
  self,
372
804
  calculation_results: CalculationResults,
@@ -375,10 +807,12 @@ class _NodeResults(_ElementResults):
375
807
  constraints: List[str],
376
808
  inputs: List[str],
377
809
  outputs: List[str],
810
+ flows: List[str],
378
811
  ):
379
812
  super().__init__(calculation_results, label, variables, constraints)
380
813
  self.inputs = inputs
381
814
  self.outputs = outputs
815
+ self.flows = flows
382
816
 
383
817
  def plot_node_balance(
384
818
  self,
@@ -386,28 +820,47 @@ class _NodeResults(_ElementResults):
386
820
  show: bool = True,
387
821
  colors: plotting.ColorType = 'viridis',
388
822
  engine: plotting.PlottingEngine = 'plotly',
823
+ scenario: Optional[Union[str, int]] = None,
824
+ mode: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
825
+ style: Literal['area', 'stacked_bar', 'line'] = 'stacked_bar',
826
+ drop_suffix: bool = True,
389
827
  ) -> Union[plotly.graph_objs.Figure, Tuple[plt.Figure, plt.Axes]]:
390
828
  """
391
829
  Plots the node balance of the Component or Bus.
392
830
  Args:
393
831
  save: Whether to save the plot or not. If a path is provided, the plot will be saved at that location.
394
832
  show: Whether to show the plot or not.
833
+ colors: The colors to use for the plot. See `flixopt.plotting.ColorType` for options.
395
834
  engine: The engine to use for plotting. Can be either 'plotly' or 'matplotlib'.
835
+ scenario: The scenario to plot. Defaults to the first scenario. Has no effect without scenarios present
836
+ mode: The mode to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
837
+ - 'flow_rate': Returns the flow_rates of the Node.
838
+ - 'flow_hours': Returns the flow_hours of the Node. [flow_hours(t) = flow_rate(t) * dt(t)]. Renames suffixes to |flow_hours.
839
+ drop_suffix: Whether to drop the suffix from the variable names.
396
840
  """
841
+ ds = self.node_balance(with_last_timestep=True, mode=mode, drop_suffix=drop_suffix)
842
+
843
+ title = f'{self.label} (flow rates)' if mode == 'flow_rate' else f'{self.label} (flow hours)'
844
+
845
+ if 'scenario' in ds.indexes:
846
+ chosen_scenario = scenario or self._calculation_results.scenarios[0]
847
+ ds = ds.sel(scenario=chosen_scenario).drop_vars('scenario')
848
+ title = f'{title} - {chosen_scenario}'
849
+
397
850
  if engine == 'plotly':
398
851
  figure_like = plotting.with_plotly(
399
- self.node_balance(with_last_timestep=True).to_dataframe(),
852
+ ds.to_dataframe(),
400
853
  colors=colors,
401
- mode='area',
402
- title=f'Flow rates of {self.label}',
854
+ style=style,
855
+ title=title,
403
856
  )
404
857
  default_filetype = '.html'
405
858
  elif engine == 'matplotlib':
406
859
  figure_like = plotting.with_matplotlib(
407
- self.node_balance(with_last_timestep=True).to_dataframe(),
860
+ ds.to_dataframe(),
408
861
  colors=colors,
409
- mode='bar',
410
- title=f'Flow rates of {self.label}',
862
+ style=style,
863
+ title=title,
411
864
  )
412
865
  default_filetype = '.png'
413
866
  else:
@@ -415,7 +868,7 @@ class _NodeResults(_ElementResults):
415
868
 
416
869
  return plotting.export_figure(
417
870
  figure_like=figure_like,
418
- default_path=self._calculation_results.folder / f'{self.label} (flow rates)',
871
+ default_path=self._calculation_results.folder / title,
419
872
  default_filetype=default_filetype,
420
873
  user_path=None if isinstance(save, bool) else pathlib.Path(save),
421
874
  show=show,
@@ -430,6 +883,7 @@ class _NodeResults(_ElementResults):
430
883
  save: Union[bool, pathlib.Path] = False,
431
884
  show: bool = True,
432
885
  engine: plotting.PlottingEngine = 'plotly',
886
+ scenario: Optional[Union[str, int]] = None,
433
887
  ) -> plotly.graph_objects.Figure:
434
888
  """
435
889
  Plots a pie chart of the flow hours of the inputs and outputs of buses or components.
@@ -441,32 +895,40 @@ class _NodeResults(_ElementResults):
441
895
  save: Whether to save the figure.
442
896
  show: Whether to show the figure.
443
897
  engine: Plotting engine to use. Only 'plotly' is implemented atm.
898
+ scenario: If scenarios are present: The scenario to plot. If None, the first scenario is used.
899
+ drop_suffix: Whether to drop the suffix from the variable names.
444
900
  """
445
- inputs = (
446
- sanitize_dataset(
447
- ds=self.solution[self.inputs],
448
- threshold=1e-5,
449
- drop_small_vars=True,
450
- zero_small_values=True,
451
- )
452
- * self._calculation_results.hours_per_timestep
901
+ inputs = sanitize_dataset(
902
+ ds=self.solution[self.inputs] * self._calculation_results.hours_per_timestep,
903
+ threshold=1e-5,
904
+ drop_small_vars=True,
905
+ zero_small_values=True,
906
+ drop_suffix='|',
453
907
  )
454
- outputs = (
455
- sanitize_dataset(
456
- ds=self.solution[self.outputs],
457
- threshold=1e-5,
458
- drop_small_vars=True,
459
- zero_small_values=True,
460
- )
461
- * self._calculation_results.hours_per_timestep
908
+ outputs = sanitize_dataset(
909
+ ds=self.solution[self.outputs] * self._calculation_results.hours_per_timestep,
910
+ threshold=1e-5,
911
+ drop_small_vars=True,
912
+ zero_small_values=True,
913
+ drop_suffix='|',
462
914
  )
915
+ inputs = inputs.sum('time')
916
+ outputs = outputs.sum('time')
917
+
918
+ title = f'{self.label} (total flow hours)'
919
+
920
+ if 'scenario' in inputs.indexes:
921
+ chosen_scenario = scenario or self._calculation_results.scenarios[0]
922
+ inputs = inputs.sel(scenario=chosen_scenario).drop_vars('scenario')
923
+ outputs = outputs.sel(scenario=chosen_scenario).drop_vars('scenario')
924
+ title = f'{title} - {chosen_scenario}'
463
925
 
464
926
  if engine == 'plotly':
465
927
  figure_like = plotting.dual_pie_with_plotly(
466
- inputs.to_dataframe().sum(),
467
- outputs.to_dataframe().sum(),
928
+ data_left=inputs.to_pandas(),
929
+ data_right=outputs.to_pandas(),
468
930
  colors=colors,
469
- title=f'Flow hours of {self.label}',
931
+ title=title,
470
932
  text_info=text_info,
471
933
  subtitles=('Inputs', 'Outputs'),
472
934
  legend_title='Flows',
@@ -476,10 +938,10 @@ class _NodeResults(_ElementResults):
476
938
  elif engine == 'matplotlib':
477
939
  logger.debug('Parameter text_info is not supported for matplotlib')
478
940
  figure_like = plotting.dual_pie_with_matplotlib(
479
- inputs.to_dataframe().sum(),
480
- outputs.to_dataframe().sum(),
941
+ data_left=inputs.to_pandas(),
942
+ data_right=outputs.to_pandas(),
481
943
  colors=colors,
482
- title=f'Total flow hours of {self.label}',
944
+ title=title,
483
945
  subtitles=('Inputs', 'Outputs'),
484
946
  legend_title='Flows',
485
947
  lower_percentage_group=lower_percentage_group,
@@ -490,7 +952,7 @@ class _NodeResults(_ElementResults):
490
952
 
491
953
  return plotting.export_figure(
492
954
  figure_like=figure_like,
493
- default_path=self._calculation_results.folder / f'{self.label} (total flow hours)',
955
+ default_path=self._calculation_results.folder / title,
494
956
  default_filetype=default_filetype,
495
957
  user_path=None if isinstance(save, bool) else pathlib.Path(save),
496
958
  show=show,
@@ -503,9 +965,25 @@ class _NodeResults(_ElementResults):
503
965
  negate_outputs: bool = False,
504
966
  threshold: Optional[float] = 1e-5,
505
967
  with_last_timestep: bool = False,
968
+ mode: Literal['flow_rate', 'flow_hours'] = 'flow_rate',
969
+ drop_suffix: bool = False,
506
970
  ) -> xr.Dataset:
507
- return sanitize_dataset(
508
- ds=self.solution[self.inputs + self.outputs],
971
+ """
972
+ Returns a dataset with the node balance of the Component or Bus.
973
+ Args:
974
+ negate_inputs: Whether to negate the input flow_rates of the Node.
975
+ negate_outputs: Whether to negate the output flow_rates of the Node.
976
+ threshold: The threshold for small values. Variables with all values below the threshold are dropped.
977
+ with_last_timestep: Whether to include the last timestep in the dataset.
978
+ mode: The mode to use for the dataset. Can be 'flow_rate' or 'flow_hours'.
979
+ - 'flow_rate': Returns the flow_rates of the Node.
980
+ - 'flow_hours': Returns the flow_hours of the Node. [flow_hours(t) = flow_rate(t) * dt(t)]. Renames suffixes to |flow_hours.
981
+ drop_suffix: Whether to drop the suffix from the variable names.
982
+ """
983
+ ds = self.solution[self.inputs + self.outputs]
984
+
985
+ ds = sanitize_dataset(
986
+ ds=ds,
509
987
  threshold=threshold,
510
988
  timesteps=self._calculation_results.timesteps_extra if with_last_timestep else None,
511
989
  negate=(
@@ -517,8 +995,15 @@ class _NodeResults(_ElementResults):
517
995
  if negate_inputs
518
996
  else None
519
997
  ),
998
+ drop_suffix='|' if drop_suffix else None,
520
999
  )
521
1000
 
1001
+ if mode == 'flow_hours':
1002
+ ds = ds * self._calculation_results.hours_per_timestep
1003
+ ds = ds.rename_vars({var: var.replace('flow_rate', 'flow_hours') for var in ds.data_vars})
1004
+
1005
+ return ds
1006
+
522
1007
 
523
1008
  class BusResults(_NodeResults):
524
1009
  """Results for a Bus"""
@@ -548,6 +1033,8 @@ class ComponentResults(_NodeResults):
548
1033
  show: bool = True,
549
1034
  colors: plotting.ColorType = 'viridis',
550
1035
  engine: plotting.PlottingEngine = 'plotly',
1036
+ style: Literal['area', 'stacked_bar', 'line'] = 'stacked_bar',
1037
+ scenario: Optional[Union[str, int]] = None,
551
1038
  ) -> plotly.graph_objs.Figure:
552
1039
  """
553
1040
  Plots the charge state of a Storage.
@@ -556,37 +1043,56 @@ class ComponentResults(_NodeResults):
556
1043
  show: Whether to show the plot or not.
557
1044
  colors: The c
558
1045
  engine: Plotting engine to use. Only 'plotly' is implemented atm.
1046
+ style: The plotting mode for the flow_rate
1047
+ scenario: The scenario to plot. Defaults to the first scenario. Has no effect without scenarios present
559
1048
 
560
1049
  Raises:
561
1050
  ValueError: If the Component is not a Storage.
562
1051
  """
563
- if engine != 'plotly':
564
- raise NotImplementedError(
565
- f'Plotting engine "{engine}" not implemented for ComponentResults.plot_charge_state.'
566
- )
567
-
568
1052
  if not self.is_storage:
569
1053
  raise ValueError(f'Cant plot charge_state. "{self.label}" is not a storage')
570
1054
 
571
- fig = plotting.with_plotly(
572
- self.node_balance(with_last_timestep=True).to_dataframe(),
573
- colors=colors,
574
- mode='area',
575
- title=f'Operation Balance of {self.label}',
576
- )
1055
+ ds = self.node_balance(with_last_timestep=True)
1056
+ charge_state = self.charge_state
1057
+
1058
+ scenario_suffix = ''
1059
+ if 'scenario' in ds.indexes:
1060
+ chosen_scenario = scenario or self._calculation_results.scenarios[0]
1061
+ ds = ds.sel(scenario=chosen_scenario).drop_vars('scenario')
1062
+ charge_state = charge_state.sel(scenario=chosen_scenario).drop_vars('scenario')
1063
+ scenario_suffix = f'--{chosen_scenario}'
1064
+ if engine == 'plotly':
1065
+ fig = plotting.with_plotly(
1066
+ ds.to_dataframe(),
1067
+ colors=colors,
1068
+ style=style,
1069
+ title=f'Operation Balance of {self.label}{scenario_suffix}',
1070
+ )
577
1071
 
578
- # TODO: Use colors for charge state?
1072
+ # TODO: Use colors for charge state?
579
1073
 
580
- charge_state = self.charge_state.to_dataframe()
581
- fig.add_trace(
582
- plotly.graph_objs.Scatter(
583
- x=charge_state.index, y=charge_state.values.flatten(), mode='lines', name=self._charge_state
1074
+ charge_state = charge_state.to_dataframe()
1075
+ fig.add_trace(
1076
+ plotly.graph_objs.Scatter(
1077
+ x=charge_state.index, y=charge_state.values.flatten(), mode='lines', name=self._charge_state
1078
+ )
1079
+ )
1080
+ elif engine=='matplotlib':
1081
+ fig, ax = plotting.with_matplotlib(
1082
+ ds.to_dataframe(),
1083
+ colors=colors,
1084
+ style=style,
1085
+ title=f'Operation Balance of {self.label}{scenario_suffix}',
584
1086
  )
585
- )
1087
+
1088
+ charge_state = charge_state.to_dataframe()
1089
+ ax.plot(charge_state.index, charge_state.values.flatten(), label=self._charge_state)
1090
+ fig.tight_layout()
1091
+ fig = fig, ax
586
1092
 
587
1093
  return plotting.export_figure(
588
1094
  fig,
589
- default_path=self._calculation_results.folder / f'{self.label} (charge state)',
1095
+ default_path=self._calculation_results.folder / f'{self.label} (charge state){scenario_suffix}',
590
1096
  default_filetype='.html',
591
1097
  user_path=None if isinstance(save, bool) else pathlib.Path(save),
592
1098
  show=show,
@@ -633,6 +1139,45 @@ class EffectResults(_ElementResults):
633
1139
  return self.solution[[name for name in self._variable_names if name.startswith(f'{element}->')]]
634
1140
 
635
1141
 
1142
+ class FlowResults(_ElementResults):
1143
+ def __init__(
1144
+ self,
1145
+ calculation_results: CalculationResults,
1146
+ label: str,
1147
+ variables: List[str],
1148
+ constraints: List[str],
1149
+ start: str,
1150
+ end: str,
1151
+ component: str,
1152
+ ):
1153
+ super().__init__(calculation_results, label, variables, constraints)
1154
+ self.start = start
1155
+ self.end = end
1156
+ self.component = component
1157
+
1158
+ @property
1159
+ def flow_rate(self) -> xr.DataArray:
1160
+ return self.solution[f'{self.label}|flow_rate']
1161
+
1162
+ @property
1163
+ def flow_hours(self) -> xr.DataArray:
1164
+ return (self.flow_rate * self._calculation_results.hours_per_timestep).rename(f'{self.label}|flow_hours')
1165
+
1166
+ @property
1167
+ def size(self) -> xr.DataArray:
1168
+ name = f'{self.label}|size'
1169
+ if name in self.solution:
1170
+ return self.solution[name]
1171
+ try:
1172
+ return DataConverter.as_dataarray(
1173
+ self._calculation_results.flow_system.flows[self.label].size,
1174
+ scenarios=self._calculation_results.scenarios
1175
+ ).rename(name)
1176
+ except _FlowSystemRestorationError:
1177
+ logger.critical(f'Size of flow {self.label}.size not availlable. Returning NaN')
1178
+ return xr.DataArray(np.nan).rename(name)
1179
+
1180
+
636
1181
  class SegmentedCalculationResults:
637
1182
  """
638
1183
  Class to store the results of a SegmentedCalculation.
@@ -824,6 +1369,7 @@ def sanitize_dataset(
824
1369
  negate: Optional[List[str]] = None,
825
1370
  drop_small_vars: bool = True,
826
1371
  zero_small_values: bool = False,
1372
+ drop_suffix: Optional[str] = None,
827
1373
  ) -> xr.Dataset:
828
1374
  """
829
1375
  Sanitizes a dataset by handling small values (dropping or zeroing) and optionally reindexing the time axis.
@@ -835,6 +1381,7 @@ def sanitize_dataset(
835
1381
  negate: The variables to negate. If None, no variables are negated.
836
1382
  drop_small_vars: If True, drops variables where all values are below threshold.
837
1383
  zero_small_values: If True, sets values below threshold to zero.
1384
+ drop_suffix: Drop suffix of data var names. Split by the provided str.
838
1385
 
839
1386
  Returns:
840
1387
  xr.Dataset: The sanitized dataset.
@@ -873,26 +1420,177 @@ def sanitize_dataset(
873
1420
  if timesteps is not None and not ds.indexes['time'].equals(timesteps):
874
1421
  ds = ds.reindex({'time': timesteps}, fill_value=np.nan)
875
1422
 
1423
+ if drop_suffix is not None:
1424
+ if not isinstance(drop_suffix, str):
1425
+ raise ValueError(f'Only pass str values to drop suffixes. Got {drop_suffix}')
1426
+ unique_dict = {}
1427
+ for var in ds.data_vars:
1428
+ new_name = var.split(drop_suffix)[0]
1429
+
1430
+ # If name already exists, keep original name
1431
+ if new_name in unique_dict.values():
1432
+ unique_dict[var] = var
1433
+ else:
1434
+ unique_dict[var] = new_name
1435
+ ds = ds.rename(unique_dict)
1436
+
876
1437
  return ds
877
1438
 
878
1439
 
879
1440
  def filter_dataset(
880
1441
  ds: xr.Dataset,
881
- variable_dims: Optional[Literal['scalar', 'time']] = None,
1442
+ variable_dims: Optional[Literal['scalar', 'time', 'scenario', 'timeonly', 'scenarioonly']] = None,
1443
+ timesteps: Optional[Union[pd.DatetimeIndex, str, pd.Timestamp]] = None,
1444
+ scenarios: Optional[Union[pd.Index, str, int]] = None,
1445
+ contains: Optional[Union[str, List[str]]] = None,
1446
+ startswith: Optional[Union[str, List[str]]] = None,
882
1447
  ) -> xr.Dataset:
883
1448
  """
884
- Filters a dataset by its dimensions.
1449
+ Filters a dataset by its dimensions, indexes, and with string filters for variable names.
885
1450
 
886
1451
  Args:
887
1452
  ds: The dataset to filter.
888
- variable_dims: The dimension of the variables to filter for.
1453
+ variable_dims: The dimension of which to get variables from.
1454
+ - 'scalar': Get scalar variables (without dimensions)
1455
+ - 'time': Get time-dependent variables (with a time dimension)
1456
+ - 'scenario': Get scenario-dependent variables (with ONLY a scenario dimension)
1457
+ - 'timeonly': Get time-dependent variables (with ONLY a time dimension)
1458
+ - 'scenarioonly': Get scenario-dependent variables (with ONLY a scenario dimension)
1459
+ timesteps: Optional time indexes to select. Can be:
1460
+ - pd.DatetimeIndex: Multiple timesteps
1461
+ - str/pd.Timestamp: Single timestep
1462
+ Defaults to all available timesteps.
1463
+ scenarios: Optional scenario indexes to select. Can be:
1464
+ - pd.Index: Multiple scenarios
1465
+ - str/int: Single scenario (int is treated as a label, not an index position)
1466
+ Defaults to all available scenarios.
1467
+ contains: Filter variables that contain this string or strings.
1468
+ If a list is provided, variables must contain ALL strings in the list.
1469
+ startswith: Filter variables that start with this string or strings.
1470
+ If a list is provided, variables must start with ANY of the strings in the list.
1471
+
1472
+ Returns:
1473
+ Filtered dataset with specified variables and indexes.
889
1474
  """
890
- if variable_dims is None:
891
- return ds
1475
+ # First filter by dimensions
1476
+ filtered_ds = ds.copy()
1477
+ if variable_dims is not None:
1478
+ if variable_dims == 'scalar':
1479
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if not filtered_ds[v].dims]]
1480
+ elif variable_dims == 'time':
1481
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if 'time' in filtered_ds[v].dims]]
1482
+ elif variable_dims == 'scenario':
1483
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if 'scenario' in filtered_ds[v].dims]]
1484
+ elif variable_dims == 'timeonly':
1485
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if filtered_ds[v].dims == ('time',)]]
1486
+ elif variable_dims == 'scenarioonly':
1487
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if filtered_ds[v].dims == ('scenario',)]]
1488
+ else:
1489
+ raise ValueError(f'Unknown variable_dims "{variable_dims}" for filter_dataset')
1490
+
1491
+ # Filter by 'contains' parameter
1492
+ if contains is not None:
1493
+ if isinstance(contains, str):
1494
+ # Single string - keep variables that contain this string
1495
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if contains in v]]
1496
+ elif isinstance(contains, list) and all(isinstance(s, str) for s in contains):
1497
+ # List of strings - keep variables that contain ALL strings in the list
1498
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if all(s in v for s in contains)]]
1499
+ else:
1500
+ raise TypeError(f"'contains' must be a string or list of strings, got {type(contains)}")
1501
+
1502
+ # Filter by 'startswith' parameter
1503
+ if startswith is not None:
1504
+ if isinstance(startswith, str):
1505
+ # Single string - keep variables that start with this string
1506
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if v.startswith(startswith)]]
1507
+ elif isinstance(startswith, list) and all(isinstance(s, str) for s in startswith):
1508
+ # List of strings - keep variables that start with ANY of the strings in the list
1509
+ filtered_ds = filtered_ds[[v for v in filtered_ds.data_vars if any(v.startswith(s) for s in startswith)]]
1510
+ else:
1511
+ raise TypeError(f"'startswith' must be a string or list of strings, got {type(startswith)}")
892
1512
 
893
- if variable_dims == 'scalar':
894
- return ds[[name for name, da in ds.data_vars.items() if len(da.dims) == 0]]
895
- elif variable_dims == 'time':
896
- return ds[[name for name, da in ds.data_vars.items() if 'time' in da.dims]]
897
- else:
898
- raise ValueError(f'Not allowed value for "filter_dataset()": {variable_dims=}')
1513
+ # Handle time selection if needed
1514
+ if timesteps is not None and 'time' in filtered_ds.dims:
1515
+ try:
1516
+ filtered_ds = filtered_ds.sel(time=timesteps)
1517
+ except KeyError as e:
1518
+ available_times = set(filtered_ds.indexes['time'])
1519
+ requested_times = set([timesteps]) if not isinstance(timesteps, pd.Index) else set(timesteps)
1520
+ missing_times = requested_times - available_times
1521
+ raise ValueError(
1522
+ f'Timesteps not found in dataset: {missing_times}. Available times: {available_times}'
1523
+ ) from e
1524
+
1525
+ # Handle scenario selection if needed
1526
+ if scenarios is not None and 'scenario' in filtered_ds.dims:
1527
+ try:
1528
+ filtered_ds = filtered_ds.sel(scenario=scenarios)
1529
+ except KeyError as e:
1530
+ available_scenarios = set(filtered_ds.indexes['scenario'])
1531
+ requested_scenarios = set([scenarios]) if not isinstance(scenarios, pd.Index) else set(scenarios)
1532
+ missing_scenarios = requested_scenarios - available_scenarios
1533
+ raise ValueError(
1534
+ f'Scenarios not found in dataset: {missing_scenarios}. Available scenarios: {available_scenarios}'
1535
+ ) from e
1536
+
1537
+ return filtered_ds
1538
+
1539
+
1540
+ def filter_dataarray_by_coord(
1541
+ da: xr.DataArray,
1542
+ **kwargs: Optional[Union[str, List[str]]]
1543
+ ) -> xr.DataArray:
1544
+ """Filter flows by node and component attributes.
1545
+
1546
+ Filters are applied in the order they are specified. All filters must match for an edge to be included.
1547
+
1548
+ To recombine filtered dataarrays, use `xr.concat`.
1549
+
1550
+ xr.concat([res.sizes(start='Fernwärme'), res.sizes(end='Fernwärme')], dim='flow')
1551
+
1552
+ Args:
1553
+ da: Flow DataArray with network metadata coordinates.
1554
+ **kwargs: Coord filters as name=value pairs.
1555
+
1556
+ Returns:
1557
+ Filtered DataArray with matching edges.
1558
+
1559
+ Raises:
1560
+ AttributeError: If required coordinates are missing.
1561
+ ValueError: If specified nodes don't exist or no matches found.
1562
+ """
1563
+ # Helper function to process filters
1564
+ def apply_filter(array, coord_name: str, coord_values: Union[Any, List[Any]]):
1565
+ # Verify coord exists
1566
+ if coord_name not in array.coords:
1567
+ raise AttributeError(f"Missing required coordinate '{coord_name}'")
1568
+
1569
+ # Convert single value to list
1570
+ val_list = [coord_values] if isinstance(coord_values, str) else coord_values
1571
+
1572
+ # Verify coord_values exist
1573
+ available = set(array[coord_name].values)
1574
+ missing = [v for v in val_list if v not in available]
1575
+ if missing:
1576
+ raise ValueError(f"{coord_name.title()} value(s) not found: {missing}")
1577
+
1578
+ # Apply filter
1579
+ return array.where(
1580
+ array[coord_name].isin(val_list) if isinstance(coord_values, list) else array[coord_name] == coord_values,
1581
+ drop=True
1582
+ )
1583
+
1584
+ # Apply filters from kwargs
1585
+ filters = {k: v for k, v in kwargs.items() if v is not None}
1586
+ try:
1587
+ for coord, values in filters.items():
1588
+ da = apply_filter(da, coord, values)
1589
+ except ValueError as e:
1590
+ raise ValueError(f"No edges match criteria: {filters}") from e
1591
+
1592
+ # Verify results exist
1593
+ if da.size == 0:
1594
+ raise ValueError(f"No edges match criteria: {filters}")
1595
+
1596
+ return da