Feature/398 feature facet plots in results

CHANGELOG.md

@@ -54,10 +54,12 @@ If upgrading from v2.x, see the [Migration Guide](https://flixopt.github.io/flix
 
 
 ### ✨ Added
+- Added faceting and animation options to plotting methods
 
 ### 💥 Breaking Changes
 
 ### ♻️ Changed
+- Changed indexer behaviour. Defaults to not indexing instead of the first value except for time. Also changed naming when indexing.
 
 ### 🗑️ Deprecated
 

flixopt/plotting.py

@@ -39,6 +39,7 @@
 import plotly.express as px
 import plotly.graph_objects as go
 import plotly.offline
+import xarray as xr
 from plotly.exceptions import PlotlyError
 
 if TYPE_CHECKING:
@@ -326,143 +327,253 @@ def process_colors(
 
 
 def with_plotly(
- data: pd.DataFrame,
+ data: pd.DataFrame | xr.DataArray | xr.Dataset,
  mode: Literal['stacked_bar', 'line', 'area', 'grouped_bar'] = 'stacked_bar',
  colors: ColorType = 'viridis',
  title: str = '',
  ylabel: str = '',
  xlabel: str = 'Time in h',
  fig: go.Figure | None = None,
+ facet_by: str | list[str] | None = None,
+ animate_by: str | None = None,
+ facet_cols: int = 3,
+ shared_yaxes: bool = True,
+ shared_xaxes: bool = True,
 ) -> go.Figure:
  """
- Plot a DataFrame with Plotly, using either stacked bars or stepped lines.
+ Plot data with Plotly using facets (subplots) and/or animation for multidimensional data.
+
+ Uses Plotly Express for convenient faceting and animation with automatic styling.
+ For simple plots without faceting, can optionally add to an existing figure.
 
  Args:
- data: A DataFrame containing the data to plot, where the index represents time (e.g., hours),
- and each column represents a separate data series.
- mode: The plotting mode. Use 'stacked_bar' for stacked bar charts, 'line' for stepped lines,
- or 'area' for stacked area charts.
- colors: Color specification, can be:
- - A string with a colorscale name (e.g., 'viridis', 'plasma')
- - A list of color strings (e.g., ['#ff0000', '#00ff00'])
- - A dictionary mapping column names to colors (e.g., {'Column1': '#ff0000'})
- title: The title of the plot.
+ data: A DataFrame or xarray DataArray/Dataset to plot.
+ mode: The plotting mode. Use 'stacked_bar' for stacked bar charts, 'line' for lines,
+ 'area' for stacked area charts, or 'grouped_bar' for grouped bar charts.
+ colors: Color specification (colormap, list, or dict mapping labels to colors).
+ title: The main title of the plot.
  ylabel: The label for the y-axis.
  xlabel: The label for the x-axis.
- fig: A Plotly figure object to plot on. If not provided, a new figure will be created.
+ fig: A Plotly figure object to plot on (only for simple plots without faceting).
+ If not provided, a new figure will be created.
+ facet_by: Dimension(s) to create facets for. Creates a subplot grid.
+ Can be a single dimension name or list of dimensions (max 2 for facet_row and facet_col).
+ If the dimension doesn't exist in the data, it will be silently ignored.
+ animate_by: Dimension to animate over. Creates animation frames.
+ If the dimension doesn't exist in the data, it will be silently ignored.
+ facet_cols: Number of columns in the facet grid (used when facet_by is single dimension).
+ shared_yaxes: Whether subplots share y-axes.
+ shared_xaxes: Whether subplots share x-axes.
 
  Returns:
- A Plotly figure object containing the generated plot.
+ A Plotly figure object containing the faceted/animated plot.
+
+ Examples:
+ Simple plot:
+
+ ```python
+ fig = with_plotly(df, mode='area', title='Energy Mix')
+ ```
+
+ Facet by scenario:
+
+ ```python
+ fig = with_plotly(ds, facet_by='scenario', facet_cols=2)
+ ```
+
+ Animate by period:
+
+ ```python
+ fig = with_plotly(ds, animate_by='period')
+ ```
+
+ Facet and animate:
+
+ ```python
+ fig = with_plotly(ds, facet_by='scenario', animate_by='period')
+ ```
  """
  if mode not in ('stacked_bar', 'line', 'area', 'grouped_bar'):
  raise ValueError(f"'mode' must be one of {{'stacked_bar','line','area', 'grouped_bar'}}, got {mode!r}")
- if data.empty:
- return go.Figure()
 
- processed_colors = ColorProcessor(engine='plotly').process_colors(colors, list(data.columns))
-
- fig = fig if fig is not None else go.Figure()
+ # Handle empty data
+ if isinstance(data, pd.DataFrame) and data.empty:
+ return go.Figure()
+ elif isinstance(data, xr.DataArray) and data.size == 0:
+ return go.Figure()
+ elif isinstance(data, xr.Dataset) and len(data.data_vars) == 0:
+ return go.Figure()
 
- if mode == 'stacked_bar':
- for i, column in enumerate(data.columns):
- fig.add_trace(
- go.Bar(
- x=data.index,
- y=data[column],
- name=column,
- marker=dict(
- color=processed_colors[i], line=dict(width=0, color='rgba(0,0,0,0)')
- ), # Transparent line with 0 width
+ # Warn if fig parameter is used with faceting
+ if fig is not None and (facet_by is not None or animate_by is not None):
+ logger.warning('The fig parameter is ignored when using faceting or animation. Creating a new figure.')
+ fig = None
+
+ # Convert xarray to long-form DataFrame for Plotly Express
+ if isinstance(data, (xr.DataArray, xr.Dataset)):
+ # Convert to long-form (tidy) DataFrame
+ # Structure: time, variable, value, scenario, period, ... (all dims as columns)
+ if isinstance(data, xr.Dataset):
+ # Stack all data variables into long format
+ df_long = data.to_dataframe().reset_index()
+ # Melt to get: time, scenario, period, ..., variable, value
+ id_vars = [dim for dim in data.dims]
+ value_vars = list(data.data_vars)
+ df_long = df_long.melt(id_vars=id_vars, value_vars=value_vars, var_name='variable', value_name='value')
+ else:
+ # DataArray
+ df_long = data.to_dataframe().reset_index()
+ if data.name:
+ df_long = df_long.rename(columns={data.name: 'value'})
+ else:
+ # Unnamed DataArray, find the value column
+ value_col = [col for col in df_long.columns if col not in data.dims][0]
+ df_long = df_long.rename(columns={value_col: 'value'})
+ df_long['variable'] = data.name or 'data'
+ else:
+ # Already a DataFrame - convert to long format for Plotly Express
+ df_long = data.reset_index()
+ if 'time' not in df_long.columns:
+ # First column is probably time
+ df_long = df_long.rename(columns={df_long.columns[0]: 'time'})
+ # Melt to long format
+ id_vars = [
+ col
+ for col in df_long.columns
+ if col in ['time', 'scenario', 'period']
+ or col in (facet_by if isinstance(facet_by, list) else [facet_by] if facet_by else [])
+ ]
+ value_vars = [col for col in df_long.columns if col not in id_vars]
+ df_long = df_long.melt(id_vars=id_vars, value_vars=value_vars, var_name='variable', value_name='value')
+
+ # Validate facet_by and animate_by dimensions exist in the data
+ available_dims = [col for col in df_long.columns if col not in ['variable', 'value']]
+
+ # Check facet_by dimensions
+ if facet_by is not None:
+ if isinstance(facet_by, str):
+ if facet_by not in available_dims:
+ logger.debug(
+ f"Dimension '{facet_by}' not found in data. Available dimensions: {available_dims}. "
+ f'Ignoring facet_by parameter.'
  )
- )
-
- fig.update_layout(
- barmode='relative',
- bargap=0, # No space between bars
- bargroupgap=0, # No space between grouped bars
+ facet_by = None
+ elif isinstance(facet_by, list):
+ # Filter out dimensions that don't exist
+ missing_dims = [dim for dim in facet_by if dim not in available_dims]
+ facet_by = [dim for dim in facet_by if dim in available_dims]
+ if missing_dims:
+ logger.debug(
+ f'Dimensions {missing_dims} not found in data. Available dimensions: {available_dims}. '
+ f'Using only existing dimensions: {facet_by if facet_by else "none"}.'
+ )
+ if len(facet_by) == 0:
+ facet_by = None
+
+ # Check animate_by dimension
+ if animate_by is not None and animate_by not in available_dims:
+ logger.debug(
+ f"Dimension '{animate_by}' not found in data. Available dimensions: {available_dims}. "
+ f'Ignoring animate_by parameter.'
  )
- if mode == 'grouped_bar':
- for i, column in enumerate(data.columns):
- fig.add_trace(go.Bar(x=data.index, y=data[column], name=column, marker=dict(color=processed_colors[i])))
+ animate_by = None
+
+ # Setup faceting parameters for Plotly Express
+ facet_row = None
+ facet_col = None
+ if facet_by:
+ if isinstance(facet_by, str):
+ # Single facet dimension - use facet_col with facet_col_wrap
+ facet_col = facet_by
+ elif len(facet_by) == 1:
+ facet_col = facet_by[0]
+ elif len(facet_by) == 2:
+ # Two facet dimensions - use facet_row and facet_col
+ facet_row = facet_by[0]
+ facet_col = facet_by[1]
+ else:
+ raise ValueError(f'facet_by can have at most 2 dimensions, got {len(facet_by)}')
+
+ # Process colors
+ all_vars = df_long['variable'].unique().tolist()
+ processed_colors = ColorProcessor(engine='plotly').process_colors(colors, all_vars)
+ color_discrete_map = {var: color for var, color in zip(all_vars, processed_colors, strict=False)}
+
+ # Create plot using Plotly Express based on mode
+ common_args = {
+ 'data_frame': df_long,
+ 'x': 'time',
+ 'y': 'value',
+ 'color': 'variable',
+ 'facet_row': facet_row,
+ 'facet_col': facet_col,
+ 'animation_frame': animate_by,
+ 'color_discrete_map': color_discrete_map,
+ 'title': title,
+ 'labels': {'value': ylabel, 'time': xlabel, 'variable': ''},
+ }
 
- fig.update_layout(
- barmode='group',
- bargap=0.2, # No space between bars
- bargroupgap=0, # space between grouped bars
- )
+ # Add facet_col_wrap for single facet dimension
+ if facet_col and not facet_row:
+ common_args['facet_col_wrap'] = facet_cols
+
+ if mode == 'stacked_bar':
+ fig = px.bar(**common_args)
+ fig.update_traces(marker_line_width=0)
+ fig.update_layout(barmode='relative', bargap=0, bargroupgap=0)
+ elif mode == 'grouped_bar':
+ fig = px.bar(**common_args)
+ fig.update_layout(barmode='group', bargap=0.2, bargroupgap=0)
  elif mode == 'line':
- for i, column in enumerate(data.columns):
- fig.add_trace(
- go.Scatter(
- x=data.index,
- y=data[column],
- mode='lines',
- name=column,
- line=dict(shape='hv', color=processed_colors[i]),
- )
- )
+ fig = px.line(**common_args, line_shape='hv') # Stepped lines
  elif mode == 'area':
- data = data.copy()
- data[(data > -1e-5) & (data < 1e-5)] = 0 # Preventing issues with plotting
- # Split columns into positive, negative, and mixed categories
- positive_columns = list(data.columns[(data >= 0).where(~np.isnan(data), True).all()])
- negative_columns = list(data.columns[(data <= 0).where(~np.isnan(data), True).all()])
- negative_columns = [column for column in negative_columns if column not in positive_columns]
- mixed_columns = list(set(data.columns) - set(positive_columns + negative_columns))
-
- if mixed_columns:
- logger.error(
- f'Data for plotting stacked lines contains columns with both positive and negative values:'
- f' {mixed_columns}. These can not be stacked, and are printed as simple lines'
- )
+ # Use Plotly Express to create the area plot (preserves animation, legends, faceting)
+ fig = px.area(**common_args, line_shape='hv')
 
- # Get color mapping for all columns
- colors_stacked = {column: processed_colors[i] for i, column in enumerate(data.columns)}
-
- for column in positive_columns + negative_columns:
- fig.add_trace(
- go.Scatter(
- x=data.index,
- y=data[column],
- mode='lines',
- name=column,
- line=dict(shape='hv', color=colors_stacked[column]),
- fill='tonexty',
- stackgroup='pos' if column in positive_columns else 'neg',
- )
- )
+ # Classify each variable based on its values
+ variable_classification = {}
+ for var in all_vars:
+ var_data = df_long[df_long['variable'] == var]['value']
+ var_data_clean = var_data[(var_data < -1e-5) | (var_data > 1e-5)]
 
- for column in mixed_columns:
- fig.add_trace(
- go.Scatter(
- x=data.index,
- y=data[column],
- mode='lines',
- name=column,
- line=dict(shape='hv', color=colors_stacked[column], dash='dash'),
+ if len(var_data_clean) == 0:
+ variable_classification[var] = 'zero'
+ else:
+ has_pos, has_neg = (var_data_clean > 0).any(), (var_data_clean < 0).any()
+ variable_classification[var] = (
+ 'mixed' if has_pos and has_neg else ('negative' if has_neg else 'positive')
  )
- )
 
- # Update layout for better aesthetics
+ # Log warning for mixed variables
+ mixed_vars = [v for v, c in variable_classification.items() if c == 'mixed']
+ if mixed_vars:
+ logger.warning(f'Variables with both positive and negative values: {mixed_vars}. Plotted as dashed lines.')
+
+ all_traces = list(fig.data)
+ for frame in fig.frames:
+ all_traces.extend(frame.data)
+
+ for trace in all_traces:
+ trace.stackgroup = variable_classification.get(trace.name, None)
+ # No opacity and no line for stacked areas
+ if trace.stackgroup is not None:
+ if hasattr(trace, 'line') and trace.line.color:
+ trace.fillcolor = trace.line.color # Will be solid by default
+ trace.line.width = 0
+
+ # Update layout with basic styling (Plotly Express handles sizing automatically)
  fig.update_layout(
- title=title,
- yaxis=dict(
- title=ylabel,
- showgrid=True, # Enable grid lines on the y-axis
- gridcolor='lightgrey', # Customize grid line color
- gridwidth=0.5, # Customize grid line width
- ),
- xaxis=dict(
- title=xlabel,
- showgrid=True, # Enable grid lines on the x-axis
- gridcolor='lightgrey', # Customize grid line color
- gridwidth=0.5, # Customize grid line width
- ),
- plot_bgcolor='rgba(0,0,0,0)', # Transparent background
- paper_bgcolor='rgba(0,0,0,0)', # Transparent paper background
- font=dict(size=14), # Increase font size for better readability
+ plot_bgcolor='rgba(0,0,0,0)',
+ paper_bgcolor='rgba(0,0,0,0)',
+ font=dict(size=12),
  )
 
+ # Update axes to share if requested (Plotly Express already handles this, but we can customize)
+ if not shared_yaxes:
+ fig.update_yaxes(matches=None)
+ if not shared_xaxes:
+ fig.update_xaxes(matches=None)
+
  return fig