"""Create a DataFrame from a Field for line plotting."""
from typing import Optional, List, Hashable, cast
import pandas as pd
from canopy.core.field import Field
from canopy.core.grid.grid_abc import Grid
[docs]
def make_lines(
field: Field,
axis: str = "time",
flatten_columns: bool = False,
layers: Optional[List[str]] = None,
) -> pd.DataFrame:
"""Create a pandas DataFrame with the field's data unstacked along the specified axis
This DataFrame is much easier to use to plot lines (e.g. with df.plot())
Parameters
----------
field : Field
The field from which to create the DataFrame
axis : str
The axis along which to unstack the data
flatten_columns : bool, optional
If True, flatten MultiIndex columns to simple strings. Default is False.
layers : Optional[List[str]], optional
Layer names to include when flatten_columns is True. Defaults to field.layers.
"""
if axis in field.grid.axis_names and field.grid.is_reduced(axis):
raise ValueError(f"Axis '{axis}' is reduced.")
if axis == "label":
raise ValueError(f"Unstacking along the 'label' axis is not allowed.")
if axis not in field.data.index.names:
raise ValueError(f"Axis '{axis}' not found in field's data.")
df = field._data.copy()
levels = cast(Hashable, [x for x in df.index.names if not x == axis])
df = cast(pd.DataFrame, df.unstack(level=levels))
if flatten_columns:
layers = layers or field.layers
df = _flatten_columns(df, layers, field.grid)
return df
def _flatten_columns(df: pd.DataFrame, layers: List[str], grid: Grid) -> pd.DataFrame:
"""Flatten MultiIndex columns to simple strings."""
if not isinstance(df.columns, pd.MultiIndex):
cols = [c for c in df.columns if c in layers]
return df[cols]
result_parts: List[pd.DataFrame] = []
for layer in layers:
layer_df = df[[layer]]
names = layer_df.columns.names
flat_cols = []
for row in layer_df.columns:
if "label" in names:
flat_cols.append(str(row[names.index("label")]))
else:
coord_axes = [ax for ax in grid.axis_names if ax in names]
values = [row[names.index(ax)] for ax in coord_axes]
flat_cols.append(
f"({', '.join(f'{v:.2f}' for v in values)})"
if len(values) > 1
else f"{values[0]:.2f}"
)
layer_df.columns = flat_cols
result_parts.append(layer_df)
if len(result_parts) == 1:
return result_parts[0]
combined = pd.concat(result_parts, axis=1, keys=layers[: len(result_parts)])
level0 = combined.columns.get_level_values(0)
level1 = combined.columns.get_level_values(1)
comb_cols = [f"{layer} - {col}" for layer, col in zip(level0, level1)]
combined.columns = comb_cols
return combined