MatplotLibAPI 3.2.13__py3-none-any.whl → 3.2.14__py3-none-any.whl

MatplotLibAPI/Timeserie.py

@@ -1,128 +1,342 @@
-# Hint for Visual Code Python Interactive window
-# %%
-from typing import Optional, Tuple
-import pandas as pd
-import matplotlib.pyplot as plt
-from matplotlib.figure import Figure
-from matplotlib.axes import Axes
-import seaborn as sns
+"""Timeserie plotting helpers."""
 
-from MatplotLibAPI.StyleTemplate import DynamicFuncFormatter, StyleTemplate, string_formatter, bmk_formatter, format_func, validate_dataframe
+from typing import Dict, Optional, Tuple, cast
 
+import matplotlib.pyplot as plt
+import pandas as pd
+import seaborn as sns
+from matplotlib.axes import Axes
+from matplotlib.figure import Figure
 
-TIMESERIE_STYLE_TEMPLATE = StyleTemplate(
-    palette='rocket',
-    format_funcs={"y": bmk_formatter, "label": string_formatter}
+from .StyleTemplate import (
+    TIMESERIE_STYLE_TEMPLATE,
+    StyleTemplate,
+    bmk_formatter,
+    format_func,
+    string_formatter,
+    validate_dataframe,
+    DynamicFuncFormatter,
+    FormatterFunc,
 )
 
+
 # region Line
 
 
-def aplot_timeserie(pd_df: pd.DataFrame,
-                      label: str,
-                      x: str,
-                      y: str,
-                      title: Optional[str] = None,
-                      style: StyleTemplate = TIMESERIE_STYLE_TEMPLATE,
-                      max_values: int = 100,
-                      sort_by: Optional[str] = None,
-                      ascending: bool = False,
-                      std: bool = False,
-                      ax: Optional[Axes] = None) -> Axes:
+def _prepare_timeserie_data(
+    pd_df: pd.DataFrame, label: str, x: str, y: str, sort_by: Optional[str]
+) -> pd.DataFrame:
+    """Prepare data for time series plotting.
 
-    validate_dataframe(pd_df, cols=[label, x, y], sort_by=sort_by)
-    style.format_funcs = format_func(style.format_funcs, label=label, x=x, y=y)
+    Parameters
+    ----------
+    pd_df : pd.DataFrame
+        Input DataFrame.
+    label : str
+        Column to group series by.
+    x : str
+        Column for x-axis values (time).
+    y : str
+        Column for y-axis values.
+    sort_by : Optional[str]
+        Column to sort by before plotting.
+
+    Returns
+    -------
+    pd.DataFrame
+        Prepared DataFrame with a datetime index.
 
-    df = pd_df[[label, x, y]].sort_values(by=[label, x])
+    Raises
+    ------
+    AttributeError
+        If required columns are missing from the DataFrame.
+    """
+    validate_dataframe(pd_df, cols=[label, x, y], sort_by=sort_by)
+    df = pd_df[[label, x, y]].sort_values(by=[label, x])  # type: ignore
     df[x] = pd.to_datetime(df[x])
     df.set_index(x, inplace=True)
+    return df
 
-    sns.set_palette(style.palette)
-    # Colors for each group
-    colors = sns.color_palette(n_colors=len(df.columns))
-    if ax is None:
-        ax = plt.gca()
 
+def _plot_timeserie_lines(
+    ax: Axes,
+    df: pd.DataFrame,
+    label: str,
+    x: str,
+    y: str,
+    std: bool,
+    style: StyleTemplate,
+    format_funcs: Optional[Dict[str, Optional[FormatterFunc]]],
+) -> None:
+    """Plot the time series lines on the axes.
+
+    Parameters
+    ----------
+    ax : Axes
+        Matplotlib axes object.
+    df : pd.DataFrame
+        DataFrame with data for plotting.
+    label : str
+        Column used to group series.
+    x : str
+        Column for x-axis values.
+    y : str
+        Column for y-axis values.
+    std : bool
+        Whether to plot rolling standard deviation.
+    style : StyleTemplate
+        Styling for the plot.
+    format_funcs : Optional[Dict[str, Optional[FormatterFunc]]]
+        Functions to format labels and values.
+    """
+    sns.set_palette(style.palette)
     # Get unique dimension_types
     label_types = df[label].unique()
-
     # Colors for each group
     colors = sns.color_palette(n_colors=len(label_types))
 
     for label_type, color in zip(label_types, colors):
-        temp_df = df[df[label] == label_type].sort_values(by=x)
+        temp_df = df[df[label] == label_type].sort_values(by=x)  # type: ignore
 
-        if style.format_funcs.get("label"):
-            label_type = style.format_funcs.get("label")(label_type)
+        if format_funcs and (label_formatter := format_funcs.get("label")):
+            label_type = label_formatter(label_type, None)
         if std:
             ma = temp_df[y].rolling(window=7, min_periods=1).mean()
             std_dev = temp_df[y].rolling(window=7, min_periods=1).std()
             # Calculate the last moving average value to include in the legend
-            last_ma_value = ma.iloc[-1]
+            last_ma_value = cast(pd.Series, ma).iloc[-1]
             # Dynamically creating the legend label
-            label_str = f"{string_formatter(label_type)} (avg 7d: {style.format_funcs[y](last_ma_value)})"
+            if format_funcs and (y_formatter := format_funcs.get(y)):
+                label_str = f"{string_formatter(label_type)} (avg 7d: {y_formatter(last_ma_value, None)})"
+            else:
+                label_str = f"{string_formatter(label_type)} (avg 7d: {last_ma_value})"
             # Plot moving average and include the last MA value in the label for the legend
-            ax.plot(temp_df.index, ma, color=color,
-                    linestyle='--', label=label_str)
+            ax.plot(temp_df.index, ma, color=color, linestyle="--", label=label_str)
 
-            ax.fill_between(temp_df.index, ma - std_dev, ma +
-                            std_dev, color=color, alpha=0.2, label='_nolegend_')
+            ax.fill_between(
+                temp_df.index,
+                ma - std_dev,
+                ma + std_dev,
+                color=color,
+                alpha=0.2,
+                label="_nolegend_",
+            )
         else:
             label_str = f"{string_formatter(label_type)}"
             # Plot moving average and include the last MA value in the label for the legend
             ax.plot(temp_df.index, temp_df[y], color=color, label=label_str)
 
+
+def _setup_timeserie_axes(
+    ax: Axes,
+    label: str,
+    x: str,
+    y: str,
+    style: StyleTemplate,
+    format_funcs: Optional[Dict[str, Optional[FormatterFunc]]],
+) -> None:
+    """Configure the axes for the time series plot.
+
+    Parameters
+    ----------
+    ax : Axes
+        Matplotlib axes object.
+    label : str
+        Column used for grouping and legend title.
+    x : str
+        Column for x-axis label.
+    y : str
+        Column for y-axis label.
+    style : StyleTemplate
+        Styling for the plot.
+    format_funcs : Optional[Dict[str, Optional[FormatterFunc]]]
+        Functions to format axis tick labels.
+    """
     ax.legend(
         title=label,
-        fontsize=style.font_size-4,
+        fontsize=style.font_size - 4,
         title_fontsize=style.font_size,
-        labelcolor='linecolor',
-        facecolor=style.background_color)
+        labelcolor="linecolor",
+        facecolor=style.background_color,
+    )
 
     ax.set_xlabel(string_formatter(x), color=style.font_color)
-    if style.format_funcs.get("x"):
-        ax.xaxis.set_major_formatter(
-            DynamicFuncFormatter(style.format_funcs.get("x")))
-    ax.tick_params(axis='x', colors=style.font_color,
-                   labelrotation=45, labelsize=style.font_size-4)
+    if format_funcs and (x_formatter := format_funcs.get("x")):
+        ax.xaxis.set_major_formatter(DynamicFuncFormatter(x_formatter))
+    ax.tick_params(
+        axis="x",
+        colors=style.font_color,
+        labelrotation=45,
+        labelsize=style.font_size - 4,
+    )
 
     ax.set_ylabel(string_formatter(y), color=style.font_color)
-    if style.format_funcs.get("y"):
-        ax.yaxis.set_major_formatter(
-            DynamicFuncFormatter(style.format_funcs.get("y")))
-    ax.tick_params(axis='y', colors=style.font_color,
-                   labelsize=style.font_size-4)
+    if format_funcs and (y_formatter := format_funcs.get("y")):
+        ax.yaxis.set_major_formatter(DynamicFuncFormatter(y_formatter))
+    ax.tick_params(axis="y", colors=style.font_color, labelsize=style.font_size - 4)
     ax.set_facecolor(style.background_color)
     ax.grid(True)
+
+
+def aplot_timeserie(
+    pd_df: pd.DataFrame,
+    label: str,
+    x: str,
+    y: str,
+    title: Optional[str] = None,
+    style: StyleTemplate = TIMESERIE_STYLE_TEMPLATE,
+    max_values: int = 100,
+    sort_by: Optional[str] = None,
+    ascending: bool = False,
+    std: bool = False,
+    ax: Optional[Axes] = None,
+) -> Axes:
+    """Plot a time series on the provided axes.
+
+    Parameters
+    ----------
+    pd_df : pd.DataFrame
+        DataFrame containing the data to plot.
+    label : str
+        Column used to group series.
+    x : str
+        Column for the x-axis values.
+    y : str
+        Column for the y-axis values.
+    title : str, optional
+        Plot title.
+    style : StyleTemplate, optional
+        Style configuration. The default is `TIMESERIE_STYLE_TEMPLATE`.
+    max_values : int, optional
+        Maximum number of rows to plot. The default is 100.
+    sort_by : str, optional
+        Column used to sort the data.
+    ascending : bool, optional
+        Sort order for the data. The default is `False`.
+    std : bool, optional
+        Whether to plot rolling standard deviation. The default is `False`.
+    ax : Axes, optional
+        Axes to draw on.
+
+    Returns
+    -------
+    Axes
+        Matplotlib axes with the time series plot.
+
+    Raises
+    ------
+    AttributeError
+        If required columns are not in the DataFrame.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> import matplotlib.pyplot as plt
+    >>> from MatplotLibAPI.Timeserie import aplot_timeserie
+    >>> data = {
+    ...     'date': pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-01', '2023-01-02']),
+    ...     'category': ['A', 'A', 'B', 'B'],
+    ...     'value': [10, 12, 15, 13]
+    ... }
+    >>> df = pd.DataFrame(data)
+    >>> fig, ax = plt.subplots()
+    >>> aplot_timeserie(df, label='category', x='date', y='value', ax=ax)
+    """
+    if ax is None:
+        ax = cast(Axes, plt.gca())
+
+    df = _prepare_timeserie_data(pd_df, label, x, y, sort_by)
+
+    format_funcs = format_func(style.format_funcs, label=label, x=x, y=y)
+
+    _plot_timeserie_lines(ax, df, label, x, y, std, style, format_funcs)
+
+    _setup_timeserie_axes(ax, label, x, y, style, format_funcs)
+
     if title:
-        ax.set_title(title, color=style.font_color, fontsize=style.font_size+4)
+        ax.set_title(title, color=style.font_color, fontsize=style.font_size + 4)
     return ax
 
 
-def fplot_timeserie(pd_df: pd.DataFrame,
-                       label: str,
-                       x: str,
-                       y: str,
-                       title: Optional[str] = None,
-                       style: StyleTemplate = TIMESERIE_STYLE_TEMPLATE,
-                       max_values: int = 100,
-                       sort_by: Optional[str] = None,
-                       ascending: bool = False,
-                       std: bool = False,
-                       figsize: Tuple[float, float] = (19.2, 10.8)) -> Figure:
-    fig = plt.figure(figsize=figsize)
+def fplot_timeserie(
+    pd_df: pd.DataFrame,
+    label: str,
+    x: str,
+    y: str,
+    title: Optional[str] = None,
+    style: StyleTemplate = TIMESERIE_STYLE_TEMPLATE,
+    max_values: int = 100,
+    sort_by: Optional[str] = None,
+    ascending: bool = False,
+    std: bool = False,
+    figsize: Tuple[float, float] = (19.2, 10.8),
+) -> Figure:
+    """Return a figure plotting the time series.
+
+    Parameters
+    ----------
+    pd_df : pd.DataFrame
+        DataFrame containing the data to plot.
+    label : str
+        Column used to group series.
+    x : str
+        Column for the x-axis values.
+    y : str
+        Column for the y-axis values.
+    title : str, optional
+        Plot title.
+    style : StyleTemplate, optional
+        Style configuration. The default is `TIMESERIE_STYLE_TEMPLATE`.
+    max_values : int, optional
+        Maximum number of rows to plot. The default is 100.
+    sort_by : str, optional
+        Column used to sort the data.
+    ascending : bool, optional
+        Sort order for the data. The default is `False`.
+    std : bool, optional
+        Whether to plot rolling standard deviation. The default is `False`.
+    figsize : tuple[float, float], optional
+        Size of the created figure. The default is (19.2, 10.8).
+
+    Returns
+    -------
+    Figure
+        Matplotlib figure containing the time series plot.
+
+    Raises
+    ------
+    AttributeError
+        If required columns are not in the DataFrame.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> from MatplotLibAPI.Timeserie import fplot_timeserie
+    >>> data = {
+    ...     'date': pd.to_datetime(['2023-01-01', '2023-01-02', '2023-01-01', '2023-01-02']),
+    ...     'category': ['A', 'A', 'B', 'B'],
+    ...     'value': [10, 12, 15, 13]
+    ... }
+    >>> df = pd.DataFrame(data)
+    >>> fig = fplot_timeserie(df, label='category', x='date', y='value')
+    """
+    fig = cast(Figure, plt.figure(figsize=figsize))
     fig.patch.set_facecolor(style.background_color)
     ax = fig.add_subplot()
-    ax = aplot_timeserie(pd_df=pd_df,
-                           label=label,
-                           x=x,
-                           y=y,
-                           title=title,
-                           style=style,
-                           max_values=max_values,
-                           std=std,
-                           sort_by=sort_by,
-                           ascending=ascending,
-                           ax=ax)
+    ax = aplot_timeserie(
+        pd_df=pd_df,
+        label=label,
+        x=x,
+        y=y,
+        title=title,
+        style=style,
+        max_values=max_values,
+        std=std,
+        sort_by=sort_by,
+        ascending=ascending,
+        ax=ax,
+    )
     return fig
+
+
 # endregion

MatplotLibAPI/Treemap.py

@@ -1,87 +1,150 @@
-# Hint for Visual Code Python Interactive window
-# %%
+"""Treemap plotting utilities."""
+
 from typing import Optional
+
 import pandas as pd
-from pandas import CategoricalDtype, BooleanDtype
 import plotly.graph_objects as go
+from pandas import BooleanDtype, CategoricalDtype
 
-from MatplotLibAPI.StyleTemplate import StyleTemplate, string_formatter, percent_formatter, validate_dataframe
+from .StyleTemplate import (
+    TREEMAP_STYLE_TEMPLATE,
+    StyleTemplate,
+    percent_formatter,
+    string_formatter,
+    validate_dataframe,
+)
 
 
-TREEMAP_STYLE_TEMPLATE = StyleTemplate(
-    background_color='black',
-    fig_border='darkgrey',
-    font_color='white',
-    palette='magma',
-    format_funcs={"y": percent_formatter,
-                  "label": string_formatter}
-)
+def aplot_treemap(
+    pd_df: pd.DataFrame,
+    path: str,
+    values: str,
+    style: StyleTemplate = TREEMAP_STYLE_TEMPLATE,
+    title: Optional[str] = None,
+    color: Optional[str] = None,
+    sort_by: Optional[str] = None,
+    ascending: bool = False,
+    max_values: int = 100,
+) -> go.Trace:
+    """Create a treemap trace from the data frame.
 
+    Parameters
+    ----------
+    pd_df : pd.DataFrame
+        DataFrame containing the data to plot.
+    path : str
+        Column representing hierarchical path.
+    values : str
+        Column containing values for each treemap block.
+    style : StyleTemplate, optional
+        Style configuration. The default is `TREEMAP_STYLE_TEMPLATE`.
+    title : str, optional
+        Plot title.
+    color : str, optional
+        Column used for coloring.
+    sort_by : str, optional
+        Column used to sort data.
+    ascending : bool, optional
+        Sort order for the data. The default is `False`.
+    max_values : int, optional
+        Maximum number of rows to plot. The default is 100.
 
-def aplot_treemap(pd_df: pd.DataFrame,
-                  path: str,
-                  values: str,
-                  style: StyleTemplate = TREEMAP_STYLE_TEMPLATE,
-                  title: Optional[str] = None,
-                  color: Optional[str] = None,
-                  sort_by: Optional[str] = None,
-                  ascending: bool = False,
-                  max_values: int = 100) -> go.Trace:
+    Returns
+    -------
+    go.Trace
+        Plotly treemap trace.
+    """
     cols = [path, values]
     if color:
         cols.append(color)
     validate_dataframe(pd_df, cols=cols, sort_by=sort_by)
     if not sort_by:
         sort_by = values
-    df = pd_df.sort_values(by=sort_by, ascending=ascending)[
-        cols].head(max_values)
-    data = {"labels": df[path],
-            "parents": [""] * len(df),
-            "values": df[values],
-            "textinfo": "label",
-            "name": title,
-            "textfont":
-                {"family": style.font_name,
-                 "size": style.font_size,
-                 "color": style.font_color}
-            }
+    df = pd_df.sort_values(by=sort_by, ascending=ascending)[cols].head(max_values)
+    data = {
+        "labels": df[path],
+        "parents": [""] * len(df),
+        "values": df[values],
+        "textinfo": "label",
+        "name": title,
+        "textfont": {
+            "family": style.font_name,
+            "size": style.font_size,
+            "color": style.font_color,
+        },
+    }
 
     if color and color in pd_df.columns:
         color_data = pd_df[color]
-        if isinstance(color_data, CategoricalDtype) or pd.api.types.is_object_dtype(color_data):
-            color_data = color_data.astype('category').cat.codes
-        elif isinstance(color_data, BooleanDtype):
+        if isinstance(
+            color_data.dtype, CategoricalDtype
+        ) or pd.api.types.is_object_dtype(color_data.dtype):
+            color_data = color_data.astype("category").cat.codes
+        elif isinstance(color_data.dtype, BooleanDtype):
             color_data = color_data.astype(int)
-        data['marker'] = dict(colorscale="Viridis",
-                              colors=color_data.to_list())
-
-    g = go.Treemap(data,
-                   root_color=style.background_color
-                   )
-
-    return g
-
-
-def fplot_treemap(pd_df: pd.DataFrame,
-                  path: str,
-                  values: str,
-                  style: StyleTemplate = TREEMAP_STYLE_TEMPLATE,
-                  title: Optional[str] = None,
-                  color: Optional[str] = None,
-                  sort_by: Optional[str] = None,
-                  ascending: bool = False,
-                  max_values: int = 100,
-                  fig: Optional[go.Figure] = None) -> go.Figure:
-
-    g = aplot_treemap(pd_df=pd_df,
-                      path=path,
-                      values=values,
-                      title=title,
-                      style=style,
-                      color=color,
-                      sort_by=sort_by,
-                      ascending=ascending,
-                      max_values=max_values)
+        data["marker"] = dict(colorscale="Viridis", colors=color_data.to_list())
+
+    g = go.Treemap(
+        data,
+        root_color=style.background_color,
+    )
+
+    return g  # type: ignore
+
+
+def fplot_treemap(
+    pd_df: pd.DataFrame,
+    path: str,
+    values: str,
+    style: StyleTemplate = TREEMAP_STYLE_TEMPLATE,
+    title: Optional[str] = None,
+    color: Optional[str] = None,
+    sort_by: Optional[str] = None,
+    ascending: bool = False,
+    max_values: int = 100,
+    fig: Optional[go.Figure] = None,
+) -> go.Figure:
+    """Return a figure containing the treemap plot.
+
+    Parameters
+    ----------
+    pd_df : pd.DataFrame
+        DataFrame containing the data to plot.
+    path : str
+        Column representing hierarchical path.
+    values : str
+        Column containing values for each treemap block.
+    style : StyleTemplate, optional
+        Style configuration. The default is `TREEMAP_STYLE_TEMPLATE`.
+    title : str, optional
+        Plot title.
+    color : str, optional
+        Column used for coloring.
+    sort_by : str, optional
+        Column used to sort data.
+    ascending : bool, optional
+        Sort order for the data. The default is `False`.
+    max_values : int, optional
+        Maximum number of rows to plot. The default is 100.
+    fig : go.Figure, optional
+        Existing figure to add the treemap to.
+
+    Returns
+    -------
+    go.Figure
+        Figure containing the treemap plot.
+    """
+    g = aplot_treemap(
+        pd_df=pd_df,
+        path=path,
+        values=values,
+        title=title,
+        style=style,
+        color=color,
+        sort_by=sort_by,
+        ascending=ascending,
+        max_values=max_values,
+    )
 
     if not fig:
         fig = go.Figure(g)
@@ -92,16 +155,11 @@ def fplot_treemap(pd_df: pd.DataFrame,
         title=title,
         plot_bgcolor=style.background_color,
         paper_bgcolor=style.background_color,
-        font=dict(
-            family=style.font_name,
-            size=style.font_size,
-            color=style.font_color
-        ),
-        showlegend=style.legend if style else True)
+        font=dict(family=style.font_name, size=style.font_size, color=style.font_color),
+        showlegend=style.legend if style else True,
+    )
 
     # Apply color scale
-    fig.update_traces(
-        marker=dict(colorscale=style.palette)
-    )
+    fig.update_traces(marker=dict(colorscale=style.palette))
 
     return fig