mgplot.growth_plot

  1"""Plot period and annual/through-the-year growth rates on the same axes.
  2
  3Key functions:
  4- calc_growth()
  5- growth_plot()
  6- series_growth_plot()
  7"""
  8
  9from collections.abc import Callable
 10from typing import NotRequired, Unpack, cast
 11
 12from matplotlib.axes import Axes
 13from numpy import nan
 14from pandas import DataFrame, Period, PeriodIndex, Series, period_range
 15
 16from mgplot.axis_utils import map_periodindex, set_labels
 17from mgplot.bar_plot import bar_plot
 18from mgplot.keyword_checking import (
 19    BaseKwargs,
 20    TransitionKwargs,
 21    package_kwargs,
 22    report_kwargs,
 23    validate_kwargs,
 24)
 25from mgplot.line_plot import line_plot
 26from mgplot.settings import DataT
 27from mgplot.utilities import check_clean_timeseries, constrain_data
 28
 29# --- constants
 30
 31# - frequency mappings
 32FREQUENCY_TO_PERIODS = {"Q": 4, "M": 12, "D": 365}
 33FREQUENCY_TO_NAME = {"Q": "Quarterly", "M": "Monthly", "D": "Daily"}
 34TWO_COLUMNS = 2
 35
 36
 37# - overarching constants
 38class GrowthKwargs(BaseKwargs):
 39    """Keyword arguments for the growth_plot function."""
 40
 41    # --- common options
 42    ax: NotRequired[Axes | None]
 43    plot_from: NotRequired[int | Period]
 44    label_series: NotRequired[bool]
 45    max_ticks: NotRequired[int]
 46    tick_relabel: NotRequired[Callable[[str], str]]
 47    # --- options passed to the line plot
 48    line_width: NotRequired[float | int]
 49    line_color: NotRequired[str]
 50    line_style: NotRequired[str]
 51    annotate_line: NotRequired[bool]
 52    line_rounding: NotRequired[bool | int]
 53    line_fontsize: NotRequired[str | int | float]
 54    line_fontname: NotRequired[str]
 55    line_anno_color: NotRequired[str]
 56    # --- options passed to the bar plot
 57    annotate_bars: NotRequired[bool]
 58    bar_fontsize: NotRequired[str | int | float]
 59    bar_fontname: NotRequired[str]
 60    bar_rounding: NotRequired[int]
 61    bar_width: NotRequired[float]
 62    bar_color: NotRequired[str]
 63    bar_anno_color: NotRequired[str]
 64    bar_rotation: NotRequired[int | float]
 65
 66
 67class SeriesGrowthKwargs(GrowthKwargs):
 68    """Keyword arguments for the series_growth_plot function."""
 69
 70    ylabel: NotRequired[str | None]
 71
 72
 73# - transition of kwargs from growth_plot to line_plot
 74common_transitions: TransitionKwargs = {
 75    # arg-to-growth_plot : (arg-to-line_plot, default_value)
 76    "label_series": ("label_series", True),
 77    "ax": ("ax", None),
 78    "max_ticks": ("max_ticks", None),
 79    "plot_from": ("plot_from", None),
 80    "report_kwargs": ("report_kwargs", None),
 81}
 82
 83to_line_plot: TransitionKwargs = common_transitions | {
 84    # arg-to-growth_plot : (arg-to-line_plot, default_value)
 85    "line_width": ("width", None),
 86    "line_color": ("color", "darkblue"),
 87    "line_style": ("style", None),
 88    "annotate_line": ("annotate", True),
 89    "line_rounding": ("rounding", None),
 90    "line_fontsize": ("fontsize", None),
 91    "line_fontname": ("fontname", None),
 92    "line_anno_color": ("annotate_color", None),
 93}
 94
 95# - constants for the bar plot
 96to_bar_plot: TransitionKwargs = common_transitions | {
 97    # arg-to-growth_plot : (arg-to-bar_plot, default_value)
 98    "bar_width": ("width", 0.8),
 99    "bar_color": ("color", "#dd0000"),
100    "annotate_bars": ("annotate", True),
101    "bar_rounding": ("rounding", None),
102    "above": ("above", False),
103    "bar_rotation": ("rotation", None),
104    "bar_fontsize": ("fontsize", None),
105    "bar_fontname": ("fontname", None),
106    "bar_anno_color": ("annotate_color", None),
107}
108
109
110# --- functions
111# - public functions
112def calc_growth(series: Series) -> DataFrame:
113    """Calculate annual and periodic growth for a pandas Series.
114
115    Args:
116        series: Series - a pandas series with a date-like PeriodIndex.
117
118    Returns:
119        DataFrame: A two column DataFrame with annual and periodic growth rates.
120
121    Raises:
122        TypeError if the series is not a pandas Series.
123        TypeError if the series index is not a PeriodIndex.
124        ValueError if the series is empty.
125        ValueError if the series index does not have a frequency of Q, M, or D.
126        ValueError if the series index has duplicates.
127
128    """
129    # --- sanity checks
130    if not isinstance(series, Series):
131        raise TypeError("The series argument must be a pandas Series")
132    if not isinstance(series.index, PeriodIndex):
133        raise TypeError("The series index must be a pandas PeriodIndex")
134    if series.empty:
135        raise ValueError("The series argument must not be empty")
136    freq = series.index.freqstr
137    if not freq or freq[0] not in FREQUENCY_TO_PERIODS:
138        raise ValueError("The series index must have a frequency of Q, M, or D")
139    if series.index.has_duplicates:
140        raise ValueError("The series index must not have duplicate values")
141
142    # --- ensure the index is complete and the date is sorted
143    complete = period_range(start=series.index.min(), end=series.index.max())
144    series = series.reindex(complete, fill_value=nan)
145    series = series.sort_index(ascending=True)
146
147    # --- calculate annual and periodic growth
148    freq = PeriodIndex(series.index).freqstr
149    if not freq or freq[0] not in FREQUENCY_TO_PERIODS:
150        raise ValueError("The series index must have a frequency of Q, M, or D")
151
152    freq_key = freq[0]
153    ppy = FREQUENCY_TO_PERIODS[freq_key]
154    annual = series.pct_change(periods=ppy) * 100
155    periodic = series.pct_change(periods=1) * 100
156    periodic_name = FREQUENCY_TO_NAME[freq_key] + " Growth"
157    return DataFrame(
158        {
159            "Annual Growth": annual,
160            periodic_name: periodic,
161        },
162    )
163
164
165def growth_plot(
166    data: DataT,
167    **kwargs: Unpack[GrowthKwargs],
168) -> Axes:
169    """Plot annual growth and periodic growth on the same axes.
170
171    Args:
172        data: A pandas DataFrame with two columns:
173        kwargs: GrowthKwargs
174
175    Returns:
176        axes: The matplotlib Axes object.
177
178    Raises:
179        TypeError if the data is not a 2-column DataFrame.
180        TypeError if the annual index is not a PeriodIndex.
181        ValueError if the annual and periodic series do not have the same index.
182
183    """
184    # --- check the kwargs
185    me = "growth_plot"
186    report_kwargs(caller=me, **kwargs)
187    validate_kwargs(GrowthKwargs, caller=me, **kwargs)
188
189    # --- data checks
190    data = check_clean_timeseries(data, me)
191    if len(data.columns) != TWO_COLUMNS:
192        raise TypeError("The data argument must be a pandas DataFrame with two columns")
193    data, kwargsd = constrain_data(data, **kwargs)
194
195    # --- get the series of interest ...
196    annual = data[data.columns[0]]
197    periodic = data[data.columns[1]]
198
199    # --- series names
200    annual.name = "Annual Growth"
201    freq = PeriodIndex(periodic.index).freqstr
202    if freq and freq[0] in FREQUENCY_TO_NAME:
203        periodic.name = FREQUENCY_TO_NAME[freq[0]] + " Growth"
204    else:
205        periodic.name = "Periodic Growth"
206
207    # --- convert PeriodIndex periodic growth data to integer indexed data.
208    saved_pi = map_periodindex(periodic)
209    if saved_pi is not None:
210        periodic = saved_pi[0]  # extract the reindexed DataFrame
211
212    # --- simple bar chart for the periodic growth
213    if "bar_anno_color" not in kwargsd or kwargsd["bar_anno_color"] is None:
214        kwargsd["bar_anno_color"] = "black" if kwargsd.get("above", False) else "white"
215    selected = package_kwargs(to_bar_plot, **kwargsd)
216    axes = bar_plot(periodic, **selected)
217
218    # --- and now the annual growth as a line
219    selected = package_kwargs(to_line_plot, **kwargsd)
220    line_plot(annual, ax=axes, **selected)
221
222    # --- fix the x-axis labels
223    if saved_pi is not None:
224        set_labels(
225            axes,
226            saved_pi[1],
227            kwargsd.get("max_ticks", 10),
228            tick_relabel=kwargsd.get("tick_relabel"),
229        )
230
231    # --- and done ...
232    return axes
233
234
235def series_growth_plot(
236    data: DataT,
237    **kwargs: Unpack[SeriesGrowthKwargs],
238) -> Axes:
239    """Plot annual and periodic growth in percentage terms from a pandas Series.
240
241    Args:
242        data: A pandas Series with an appropriate PeriodIndex.
243        kwargs: SeriesGrowthKwargs
244
245    """
246    # --- check the kwargs
247    me = "series_growth_plot"
248    report_kwargs(caller=me, **kwargs)
249    validate_kwargs(SeriesGrowthKwargs, caller=me, **kwargs)
250
251    # --- sanity checks
252    if not isinstance(data, Series):
253        raise TypeError("The data argument to series_growth_plot() must be a pandas Series")
254
255    # --- calculate growth and plot - add ylabel
256    ylabel: str | None = kwargs.pop("ylabel", None)
257    if ylabel is not None:
258        print(f"Did you intend to specify a value for the 'ylabel' in {me}()?")
259    ylabel = "Growth (%)" if ylabel is None else ylabel
260    growth = calc_growth(data)
261    ax = growth_plot(growth, **cast("GrowthKwargs", kwargs))
262    ax.set_ylabel(ylabel)
263    return ax