mgplot

40class BarKwargs(BaseKwargs):
41    """Keyword arguments for the bar_plot function."""
42
43    # --- options for the entire bar plot
44    ax: NotRequired[Axes | None]
45    stacked: NotRequired[bool]
46    horizontal: NotRequired[bool]
47    max_ticks: NotRequired[int]
48    tick_relabel: NotRequired[Callable[[str], str]]
49    plot_from: NotRequired[int | Period]
50    label_rotation: NotRequired[int | float]
51    # --- options for each bar ...
52    color: NotRequired[str | Sequence[str]]
53    label_series: NotRequired[bool | Sequence[bool]]
54    width: NotRequired[float | int | Sequence[float | int]]
55    zorder: NotRequired[int | float | Sequence[int | float]]
56    # --- options for bar annotations
57    annotate: NotRequired[bool]
58    fontsize: NotRequired[int | float | str]
59    fontname: NotRequired[str]
60    rounding: NotRequired[int]
61    rotation: NotRequired[int | float]
62    annotate_color: NotRequired[str]
63    above: NotRequired[bool]

22class FillBetweenKwargs(BaseKwargs):
23    """Keyword arguments for the fill_between_plot function."""
24
25    ax: NotRequired[Axes | None]
26    color: NotRequired[str]
27    alpha: NotRequired[float]
28    label: NotRequired[str | None]
29    linewidth: NotRequired[float]
30    edgecolor: NotRequired[str | None]
31    zorder: NotRequired[int | float]
32    plot_from: NotRequired[int | None]
33    max_ticks: NotRequired[int]
34    tick_relabel: NotRequired[Callable[[str], str]]

34class FinaliseKwargs(BaseKwargs):
35    """Keyword arguments for the finalise_plot function."""
36
37    # --- value options
38    suptitle: NotRequired[str | None]
39    title: NotRequired[str | None]
40    xlabel: NotRequired[str | None]
41    ylabel: NotRequired[str | None]
42    xlim: NotRequired[tuple[float, float] | None]
43    ylim: NotRequired[tuple[float, float] | None]
44    xticks: NotRequired[list[float] | None]
45    yticks: NotRequired[list[float] | None]
46    xscale: NotRequired[str | None]
47    yscale: NotRequired[str | None]
48    # --- splat options
49    legend: NotRequired[bool | dict[str, Any] | None]
50    axhspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
51    axvspan: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
52    axhline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
53    axvline: NotRequired[dict[str, Any] | Sequence[dict[str, Any]] | None]
54    # --- options for annotations
55    lfooter: NotRequired[str]
56    rfooter: NotRequired[str]
57    lheader: NotRequired[str]
58    rheader: NotRequired[str]
59    # --- file/save options
60    pre_tag: NotRequired[str]
61    tag: NotRequired[str]
62    filename: NotRequired[str]
63    chart_dir: NotRequired[str]
64    file_type: NotRequired[str]
65    dpi: NotRequired[int]
66    figsize: NotRequired[tuple[float, float]]
67    show: NotRequired[bool]
68    # --- other options
69    preserve_lims: NotRequired[bool]
70    remove_legend: NotRequired[bool]
71    zero_y: NotRequired[bool]
72    y0: NotRequired[bool]
73    x0: NotRequired[bool]
74    axisbelow: NotRequired[bool]
75    dont_save: NotRequired[bool]
76    dont_close: NotRequired[bool]
77    axes_only: NotRequired[bool]

39class GrowthKwargs(BaseKwargs):
40    """Keyword arguments for the growth_plot function."""
41
42    # --- common options
43    ax: NotRequired[Axes | None]
44    plot_from: NotRequired[int | Period]
45    label_series: NotRequired[bool]
46    max_ticks: NotRequired[int]
47    tick_relabel: NotRequired[Callable[[str], str]]
48    # --- options passed to the line plot
49    line_width: NotRequired[float | int]
50    line_color: NotRequired[str]
51    line_style: NotRequired[str]
52    annotate_line: NotRequired[bool]
53    line_rounding: NotRequired[bool | int]
54    line_fontsize: NotRequired[str | int | float]
55    line_fontname: NotRequired[str]
56    line_anno_color: NotRequired[str]
57    # --- options passed to the bar plot
58    annotate_bars: NotRequired[bool]
59    bar_fontsize: NotRequired[str | int | float]
60    bar_fontname: NotRequired[str]
61    bar_rounding: NotRequired[int]
62    bar_width: NotRequired[float]
63    bar_color: NotRequired[str]
64    bar_anno_color: NotRequired[str]
65    bar_rotation: NotRequired[int | float]

34class LineKwargs(BaseKwargs):
35    """Keyword arguments for the line_plot function."""
36
37    # --- options for the entire line plot
38    ax: NotRequired[Axes | None]
39    style: NotRequired[str | Sequence[str]]
40    width: NotRequired[float | int | Sequence[float | int]]
41    color: NotRequired[str | Sequence[str]]
42    alpha: NotRequired[float | Sequence[float]]
43    drawstyle: NotRequired[str | Sequence[str] | None]
44    marker: NotRequired[str | Sequence[str] | None]
45    markersize: NotRequired[float | Sequence[float] | int | None]
46    zorder: NotRequired[int | float | Sequence[int | float]]
47    dropna: NotRequired[bool | Sequence[bool]]
48    annotate: NotRequired[bool | Sequence[bool]]
49    rounding: NotRequired[Sequence[int | bool] | int | bool | None]
50    fontsize: NotRequired[Sequence[str | int | float] | str | int | float]
51    fontname: NotRequired[str | Sequence[str]]
52    rotation: NotRequired[Sequence[int | float] | int | float]
53    annotate_color: NotRequired[str | Sequence[str] | bool | Sequence[bool] | None]
54    near_end: NotRequired[float]
55    plot_from: NotRequired[int | Period | None]
56    label_series: NotRequired[bool | Sequence[bool] | None]
57    max_ticks: NotRequired[int]
58    tick_relabel: NotRequired[Callable[[str], str]]

30class PostcovidKwargs(LineKwargs):
31    """Keyword arguments for the post-COVID plot."""
32
33    start_r: NotRequired[Period]  # start of regression period
34    end_r: NotRequired[Period]  # end of regression period

32class RunKwargs(LineKwargs):
33    """Keyword arguments for the run_plot function."""
34
35    threshold: NotRequired[float]
36    direction: NotRequired[str]
37    highlight_color: NotRequired[str | Sequence[str]]
38    highlight_label: NotRequired[str | Sequence[str]]

68class SeriesGrowthKwargs(GrowthKwargs):
69    """Keyword arguments for the series_growth_plot function."""
70
71    ylabel: NotRequired[str | None]

41class SummaryKwargs(BaseKwargs):
42    """Keyword arguments for the summary_plot function."""
43
44    ax: NotRequired[Axes | None]
45    verbose: NotRequired[bool]
46    middle: NotRequired[float]
47    plot_type: NotRequired[str]
48    plot_from: NotRequired[int | Period]
49    legend: NotRequired[bool | dict[str, Any] | None]
50    xlabel: NotRequired[str | None]

158def abbreviate_state(state: str) -> str:
159    """Abbreviate long-form state names.
160
161    Args:
162        state: str - the long-form state name.
163
164    Return the abbreviation for a state name.
165
166    """
167    return _state_names_multi.get(state.lower(), state)

252def bar_plot(data: DataT, **kwargs: Unpack[BarKwargs]) -> Axes:
253    """Create a bar plot from the given data.
254
255    Each column in the DataFrame will be stacked on top of each other,
256    with positive values above zero and negative values below zero.
257
258    Args:
259        data: Series | DataFrame - The data to plot. Can be a DataFrame or a Series.
260        **kwargs: BarKwargs - Additional keyword arguments for customization.
261        (see BarKwargs for details)
262
263    Note: This function does not assume all data is timeseries with a PeriodIndex.
264
265    Returns:
266        axes: Axes - The axes for the plot.
267
268    """
269    # --- check the kwargs
270    report_kwargs(caller=ME, **kwargs)
271    validate_kwargs(schema=BarKwargs, caller=ME, **kwargs)
272
273    # --- get the data
274    # no call to check_clean_timeseries here, as bar plots are not
275    # necessarily timeseries data. If the data is a Series, it will be
276    # converted to a DataFrame with a single column.
277    df = DataFrame(data)  # really we are only plotting DataFrames
278    df, kwargs_d = constrain_data(df, **kwargs)
279    item_count = len(df.columns)
280
281    # --- deal with string indices
282    saved_strings = map_stringindex(df)
283    if saved_strings is not None:
284        df = saved_strings[0]
285
286    # --- deal with complete PeriodIndex indices
287    saved_pi = map_periodindex(df)
288    if saved_pi is not None:
289        df = saved_pi[0]  # extract the reindexed DataFrame from the PeriodIndex
290
291    # --- set up the default arguments
292    chart_defaults: dict[str, bool | int] = {
293        "stacked": False,
294        "horizontal": False,
295        "max_ticks": DEFAULT_MAX_TICKS,
296        "label_series": item_count > 1,
297        "label_rotation": 0,
298    }
299    chart_args = {k: kwargs_d.get(k, v) for k, v in chart_defaults.items()}
300
301    # --- horizontal bars are for categorical data, not PeriodIndex timeseries
302    horizontal = bool(chart_args["horizontal"])
303    if horizontal and saved_pi is not None:
304        print(f"Warning: horizontal=True is not supported with a PeriodIndex in {ME}(); plotting vertical.")
305        horizontal = False
306
307    # --- single series + one colour per bar => per-bar colours
308    user_color = kwargs_d.get("color")
309    if (
310        item_count == 1
311        and isinstance(user_color, Sequence)
312        and not isinstance(user_color, str)
313        and len(user_color) == len(df)
314    ):
315        kwargs_d["color"] = [list(user_color)]  # one series whose colour is a per-bar list
316
317    bar_defaults = {
318        "color": get_color_list(item_count),
319        "width": get_setting("bar_width"),
320        "label_series": item_count > 1,
321        "zorder": None,
322    }
323    above = kwargs_d.get("above", False)
324    anno_args: AnnoKwargs = {
325        "annotate": kwargs_d.get("annotate", False),
326        "fontsize": kwargs_d.get("fontsize", "small"),
327        "fontname": kwargs_d.get("fontname", "Helvetica"),
328        "rotation": kwargs_d.get("rotation", 0),
329        "rounding": kwargs_d.get("rounding", True),
330        "color": kwargs_d.get("annotate_color", "black" if above else "white"),
331        "above": above,
332    }
333    bar_args, remaining_kwargs = apply_defaults(item_count, bar_defaults, kwargs_d)
334
335    # --- plot the data
336    axes, remaining_kwargs = get_axes(**dict(remaining_kwargs))
337    if chart_args["stacked"]:
338        stacked(axes, df, anno_args, horizontal=horizontal, **bar_args)
339    else:
340        grouped(axes, df, anno_args, horizontal=horizontal, **bar_args)
341
342    # --- handle index labels and rotation
343    if saved_strings is not None:
344        if horizontal:
345            axes.set_yticks(range(len(saved_strings[1])))
346            axes.set_yticklabels(saved_strings[1], rotation=chart_args["label_rotation"])
347        else:
348            axes.set_xticks(range(len(saved_strings[1])))
349            axes.set_xticklabels(saved_strings[1], rotation=chart_args["label_rotation"])
350    elif saved_pi is not None:
351        set_labels(
352            axes,
353            saved_pi[1],
354            chart_args["max_ticks"],
355            rotation=chart_args["label_rotation"],
356            tick_relabel=kwargs_d.get("tick_relabel"),
357        )
358
359    return axes

138def bar_plot_finalise(
139    data: DataT,
140    **kwargs: Unpack[BPFKwargs],
141) -> None:
142    """Call bar_plot() and finalise_plot().
143
144    Args:
145        data: The data to be plotted.
146        kwargs: Combined bar plot and finalise plot keyword arguments.
147
148    """
149    validate_kwargs(schema=BPFKwargs, caller="bar_plot_finalise", **kwargs)
150    kwargs = impose_legend(kwargs=kwargs, data=data)
151    plot_then_finalise(
152        data,
153        function=bar_plot,
154        **kwargs,
155    )

113def calc_growth(series: Series) -> DataFrame:
114    """Calculate annual and periodic growth for a pandas Series.
115
116    Args:
117        series: Series - a pandas series with a date-like PeriodIndex.
118
119    Returns:
120        DataFrame: A two column DataFrame with annual and periodic growth rates.
121
122    Raises:
123        TypeError if the series is not a pandas Series.
124        TypeError if the series index is not a PeriodIndex.
125        ValueError if the series is empty.
126        ValueError if the series index does not have a frequency of Q, M, or D.
127        ValueError if the series index has duplicates.
128
129    """
130    # --- sanity checks
131    if not isinstance(series, Series):
132        raise TypeError("The series argument must be a pandas Series")
133    if not isinstance(series.index, PeriodIndex):
134        raise TypeError("The series index must be a pandas PeriodIndex")
135    if series.empty:
136        raise ValueError("The series argument must not be empty")
137    freq = series.index.freqstr
138    if not freq or freq[0] not in FREQUENCY_TO_PERIODS:
139        raise ValueError("The series index must have a frequency of Q, M, or D")
140    if series.index.has_duplicates:
141        raise ValueError("The series index must not have duplicate values")
142
143    # --- ensure the index is complete and the date is sorted
144    complete = period_range(start=series.index.min(), end=series.index.max())
145    series = series.reindex(complete, fill_value=nan)
146    series = series.sort_index(ascending=True)
147
148    # --- calculate annual and periodic growth
149    freq = PeriodIndex(series.index).freqstr
150    if not freq or freq[0] not in FREQUENCY_TO_PERIODS:
151        raise ValueError("The series index must have a frequency of Q, M, or D")
152
153    freq_key = freq[0]
154    ppy = FREQUENCY_TO_PERIODS[freq_key]
155    annual = series.pct_change(periods=ppy) * 100
156    periodic = series.pct_change(periods=1) * 100
157    periodic_name = FREQUENCY_TO_NAME[freq_key] + " Growth"
158    return DataFrame(
159        {
160            "Annual Growth": annual,
161            periodic_name: periodic,
162        },
163    )

176@contextmanager
177def chart_subdir(name: str, *, clear: bool = False) -> Generator[str]:
178    """Temporarily redirect chart output to a subdirectory of the current chart_dir.
179
180    Args:
181        name: str - subdirectory name (relative to the current chart_dir).
182        clear: bool - if True, clear the subdirectory of graph-image files
183            on entry. Only set this on the first use of a subdirectory in a
184            notebook; later cells writing to the same subdirectory should
185            leave it False.
186
187    Yields:
188        str - the subdirectory path.
189
190    Note: The previous chart directory is restored on exit, even if an
191    exception is raised.
192
193    """
194    main_dir = get_setting("chart_dir")
195    sub_dir = str(Path(main_dir) / name)
196    set_chart_dir(sub_dir)
197    if clear:
198        clear_chart_dir()
199    try:
200        yield sub_dir
201    finally:
202        set_setting("chart_dir", main_dir)

148def clear_chart_dir() -> None:
149    """Remove all graph-image files from the global chart_dir."""
150    chart_dir = get_setting("chart_dir")
151    Path(chart_dir).mkdir(parents=True, exist_ok=True)
152    for ext in IMAGE_EXTENSIONS:
153        for fs_object in Path(chart_dir).glob(f"*.{ext}"):
154            if fs_object.is_file():
155                fs_object.unlink()

103def colorise_list(party_list: Iterable[str]) -> list[str]:
104    """Return a list of party/state colors for a party_list."""
105    return [get_color(x) for x in party_list]

108def contrast(orig_color: str) -> str:
109    """Provide a contrasting color to any party color."""
110    new_color = DEFAULT_CONTRAST_COLOR
111    match orig_color:
112        case "royalblue":
113            new_color = "indianred"
114        case "indianred":
115            new_color = "royalblue"
116
117        case "darkorange":
118            new_color = "mediumblue"
119        case "mediumblue":
120            new_color = "darkorange"
121
122        case "seagreen":
123            new_color = "darkblue"
124
125        case color if color == DEFAULT_UNKNOWN_COLOR:
126            new_color = "hotpink"
127
128    return new_color

 37def fill_between_plot(data: DataFrame, **kwargs: Unpack[FillBetweenKwargs]) -> Axes:
 38    """Plot a filled region between lower and upper bounds.
 39
 40    Args:
 41        data: DataFrame - A two-column DataFrame with PeriodIndex.
 42              The first column is the lower bound, the second is the upper bound.
 43        kwargs: FillBetweenKwargs - keyword arguments for the plot.
 44
 45    Returns:
 46        Axes - matplotlib Axes object.
 47
 48    Raises:
 49        TypeError: If data is not a DataFrame.
 50        ValueError: If data does not have exactly two columns.
 51
 52    """
 53    # --- validate inputs
 54    report_kwargs(caller=ME, **kwargs)
 55    validate_kwargs(schema=FillBetweenKwargs, caller=ME, **kwargs)
 56
 57    if not isinstance(data, DataFrame):
 58        raise TypeError(f"data must be a DataFrame for {ME}()")
 59
 60    if len(data.columns) != REQUIRED_COLUMNS:
 61        raise ValueError(f"data must have exactly two columns for {ME}(), got {len(data.columns)}")
 62
 63    # --- check and constrain data
 64    data = check_clean_timeseries(data, ME)
 65    data, kwargs_d = constrain_data(data, **kwargs)
 66
 67    # --- handle PeriodIndex conversion
 68    saved_pi = map_periodindex(data)
 69    if saved_pi is not None:
 70        data = saved_pi[0]
 71
 72    # --- get axes
 73    axes, kwargs_d = get_axes(**kwargs_d)
 74
 75    if data.empty or data.isna().all().all():
 76        print(f"Warning: No data to plot in {ME}().")
 77        return axes
 78
 79    # --- extract bounds
 80    lower = data.iloc[:, 0]
 81    upper = data.iloc[:, 1]
 82
 83    # --- extract plot arguments
 84    color = kwargs_d.get("color", DEFAULT_COLOR)
 85    alpha = kwargs_d.get("alpha", DEFAULT_ALPHA)
 86    label = kwargs_d.get("label", None)
 87    linewidth = kwargs_d.get("linewidth", 0)
 88    edgecolor = kwargs_d.get("edgecolor", None)
 89    zorder = kwargs_d.get("zorder", None)
 90
 91    # --- plot
 92    axes.fill_between(
 93        data.index,
 94        lower,
 95        upper,
 96        color=color,
 97        alpha=alpha,
 98        label=label,
 99        linewidth=linewidth,
100        edgecolor=edgecolor,
101        zorder=zorder,
102    )
103
104    # --- set axis labels
105    if saved_pi is not None:
106        set_labels(
107            axes,
108            saved_pi[1],
109            kwargs_d.get("max_ticks", get_setting("max_ticks")),
110            tick_relabel=kwargs_d.get("tick_relabel"),
111        )
112
113    return axes

158def fill_between_plot_finalise(
159    data: DataFrame,
160    **kwargs: Unpack[FBPFKwargs],
161) -> None:
162    """Call fill_between_plot() and finalise_plot().
163
164    Args:
165        data: DataFrame with two columns (lower bound, upper bound).
166        kwargs: Combined fill_between plot and finalise plot keyword arguments.
167
168    """
169    validate_kwargs(schema=FBPFKwargs, caller="fill_between_plot_finalise", **kwargs)
170    kwargs = impose_legend(kwargs=kwargs, data=data)
171    plot_then_finalise(
172        data,
173        function=fill_between_plot,
174        **kwargs,
175    )

408def finalise_plot(axes: Axes, **kwargs: Unpack[FinaliseKwargs]) -> None:
409    """Finalise and save plots to the file system.
410
411    The filename for the saved plot is constructed from the global
412    chart_dir, the plot's title, any specified tag text, and the
413    file_type for the plot.
414
415    Args:
416        axes: Axes - matplotlib axes object - required
417        kwargs: FinaliseKwargs
418
419    """
420    # --- check the kwargs
421    report_kwargs(caller=ME, **kwargs)
422    validate_kwargs(schema=FinaliseKwargs, caller=ME, **kwargs)
423
424    # --- sanity checks
425    if len(axes.get_children()) < 1:
426        print(f"Warning: {ME}() called with an empty axes, which was ignored.")
427        return
428
429    # --- remember axis-limits should we need to restore thems
430    xlim, ylim = axes.get_xlim(), axes.get_ylim()
431
432    # margins
433    axes.margins(DEFAULT_MARGIN)
434    axes.autoscale(tight=False)  # This is problematic ...
435
436    apply_kwargs(axes, **kwargs)
437
438    # tight layout and save the figure
439    fig = axes.figure
440    axes_only = kwargs.get("axes_only", False)
441    if not axes_only and (suptitle := kwargs.get("suptitle")):
442        fig.suptitle(suptitle)
443    if kwargs.get("preserve_lims"):
444        # restore the original limits of the axes
445        axes.set_xlim(xlim)
446        axes.set_ylim(ylim)
447    if not axes_only and not isinstance(fig, SubFigure):
448        fig.tight_layout(pad=TIGHT_LAYOUT_PAD)
449    apply_late_kwargs(axes, **kwargs)
450    # axvspan/axvline in late_kwargs may have widened xlim beyond what
451    # set_labels() last saw; regenerate ticks from the updated view.
452    refresh_period_labels(axes)
453    # de-collide end-of-line annotations now that the layout/limits are final
454    resolve_annotation_collisions(axes)
455    legend = axes.get_legend()
456    if legend and kwargs.get("remove_legend", False):
457        legend.remove()
458    if not axes_only and not isinstance(fig, SubFigure):
459        save_to_file(fig, **kwargs)
460
461    # show the plot in Jupyter Lab
462    if not axes_only and kwargs.get("show"):
463        plt.show()
464
465    # And close - the figure this axes belongs to, not pyplot's current figure
466    if not axes_only and not kwargs.get("dont_close", False):
467        root = fig
468        while isinstance(root, SubFigure):
469            root = root.figure
470        plt.close(root)

 45def get_color(s: str) -> str:
 46    """Return a matplotlib color for a party label or an Australian state/territory.
 47
 48    Args:
 49        s: str - the party label or Australian state/territory name.
 50
 51    Returns a color string that can be used in matplotlib plots.
 52
 53    """
 54    # Flattened color map for better readability
 55    color_map: dict[str, str] = {
 56        # --- Australian states and territories
 57        "wa": "gold",
 58        "western australia": "gold",
 59        "sa": "red",
 60        "south australia": "red",
 61        "nt": "#CC7722",  # ochre
 62        "northern territory": "#CC7722",
 63        "nsw": "deepskyblue",
 64        "new south wales": "deepskyblue",
 65        "act": "blue",
 66        "australian capital territory": "blue",
 67        "vic": "navy",
 68        "victoria": "navy",
 69        "tas": "seagreen",  # bottle green #006A4E?
 70        "tasmania": "seagreen",
 71        "qld": "#c32148",  # a lighter maroon
 72        "queensland": "#c32148",
 73        "australia": "grey",
 74        "aus": "grey",
 75        # --- political parties
 76        "dissatisfied": "darkorange",  # must be before satisfied
 77        "satisfied": "mediumblue",
 78        "lnp": "royalblue",
 79        "l/np": "royalblue",
 80        "liberal": "royalblue",
 81        "liberals": "royalblue",
 82        "coalition": "royalblue",
 83        "dutton": "royalblue",
 84        "ley": "royalblue",
 85        "liberal and/or nationals": "royalblue",
 86        "nat": "forestgreen",
 87        "nats": "forestgreen",
 88        "national": "forestgreen",
 89        "nationals": "forestgreen",
 90        "alp": "#dd0000",
 91        "labor": "#dd0000",
 92        "albanese": "#dd0000",
 93        "grn": "limegreen",
 94        "green": "limegreen",
 95        "greens": "limegreen",
 96        "other": "darkorange",
 97        "oth": "darkorange",
 98    }
 99
100    return color_map.get(s.lower(), DEFAULT_UNKNOWN_COLOR)

21def get_party_palette(party_text: str) -> str:
22    """Return a matplotlib color-map name based on party_text.
23
24    Works for Australian major political parties.
25
26    Args:
27        party_text: str - the party label or name.
28
29    """
30    # Note: light to dark colormaps work best for sequential data visualization
31    match party_text.lower():
32        case "alp" | "labor":
33            return "Reds"
34        case "l/np" | "coalition":
35            return "Blues"
36        case "grn" | "green" | "greens":
37            return "Greens"
38        case "oth" | "other":
39            return "YlOrBr"
40        case "onp" | "one nation":
41            return "YlGnBu"
42    return DEFAULT_PARTY_PALETTE

104def get_setting(setting: str) -> Any:
105    """Get a setting from the global settings.
106
107    Args:
108        setting: str - name of the setting to get.
109
110    Raises:
111        KeyError: if the setting is not found
112
113    Returns:
114        value: Any - the value of the setting
115
116    """
117    if setting not in get_fields():
118        raise KeyError(f"Setting '{setting}' not found in mgplot_defaults.")
119    return getattr(mgplot_defaults, setting)

166def growth_plot(
167    data: DataT,
168    **kwargs: Unpack[GrowthKwargs],
169) -> Axes:
170    """Plot annual growth and periodic growth on the same axes.
171
172    Args:
173        data: A pandas DataFrame with two columns:
174        kwargs: GrowthKwargs
175
176    Returns:
177        axes: The matplotlib Axes object.
178
179    Raises:
180        TypeError if the data is not a 2-column DataFrame.
181        TypeError if the annual index is not a PeriodIndex.
182        ValueError if the annual and periodic series do not have the same index.
183
184    """
185    # --- check the kwargs
186    me = "growth_plot"
187    report_kwargs(caller=me, **kwargs)
188    validate_kwargs(GrowthKwargs, caller=me, **kwargs)
189
190    # --- data checks
191    data = check_clean_timeseries(data, me)
192    if len(data.columns) != TWO_COLUMNS:
193        raise TypeError("The data argument must be a pandas DataFrame with two columns")
194    data, kwargsd = constrain_data(data, **kwargs)
195
196    # --- get the series of interest ...
197    annual = data[data.columns[0]]
198    periodic = data[data.columns[1]]
199
200    # --- series names
201    annual.name = "Annual Growth"
202    freq = PeriodIndex(periodic.index).freqstr
203    if freq and freq[0] in FREQUENCY_TO_NAME:
204        periodic.name = FREQUENCY_TO_NAME[freq[0]] + " Growth"
205    else:
206        periodic.name = "Periodic Growth"
207
208    # --- convert PeriodIndex periodic growth data to integer indexed data.
209    saved_pi = map_periodindex(periodic)
210    if saved_pi is not None:
211        periodic = saved_pi[0]  # extract the reindexed DataFrame
212
213    # --- simple bar chart for the periodic growth
214    if "bar_anno_color" not in kwargsd or kwargsd["bar_anno_color"] is None:
215        kwargsd["bar_anno_color"] = "black" if kwargsd.get("above", False) else "white"
216    selected = package_kwargs(to_bar_plot, **kwargsd)
217    axes = bar_plot(periodic, **selected)
218
219    # --- and now the annual growth as a line
220    selected = package_kwargs(to_line_plot, **kwargsd)
221    line_plot(annual, ax=axes, **selected)
222
223    # --- fix the x-axis labels
224    if saved_pi is not None:
225        set_labels(
226            axes,
227            saved_pi[1],
228            kwargsd.get("max_ticks", 10),
229            tick_relabel=kwargsd.get("tick_relabel"),
230        )
231
232    # --- and done ...
233    return axes

178def growth_plot_finalise(data: DataT, **kwargs: Unpack[GrowthPFKwargs]) -> None:
179    """Call growth_plot() and finalise_plot().
180
181    Args:
182        data: The growth data to be plotted.
183        kwargs: Combined growth plot and finalise plot keyword arguments.
184
185    Note:
186        Use this when you are providing the raw growth data. Don't forget to
187        set the ylabel in kwargs.
188
189    """
190    validate_kwargs(schema=GrowthPFKwargs, caller="growth_plot_finalise", **kwargs)
191    kwargs = impose_legend(kwargs=kwargs, force=True)
192    plot_then_finalise(data=data, function=growth_plot, **kwargs)

161def line_plot(data: DataT, **kwargs: Unpack[LineKwargs]) -> Axes:
162    """Build a single or multi-line plot.
163
164    Args:
165        data: DataFrame | Series - data to plot
166        kwargs: LineKwargs - keyword arguments for the line plot
167
168    Returns:
169    - axes: Axes - the axes object for the plot
170
171    """
172    # --- check the kwargs
173    report_kwargs(caller=ME, **kwargs)
174    validate_kwargs(schema=LineKwargs, caller=ME, **kwargs)
175
176    # --- check the data
177    data = check_clean_timeseries(data, ME)
178    df = DataFrame(data)  # we are only plotting DataFrames
179    df, kwargs_d = constrain_data(df, **kwargs)
180
181    # --- convert PeriodIndex to Integer Index
182    saved_pi = map_periodindex(df)
183    if saved_pi is not None:
184        df = saved_pi[0]
185
186    if isinstance(df.index, PeriodIndex):
187        print("Internal error: data is still a PeriodIndex - come back here and fix it")
188
189    # --- Let's plot
190    axes, kwargs_d = get_axes(**kwargs_d)  # get the axes to plot on
191    if df.empty or df.isna().all().all():
192        # Note: finalise plot should ignore an empty axes object
193        print(f"Warning: No data to plot in {ME}().")
194        return axes
195
196    # --- get the arguments for each line we will plot ...
197    item_count = len(df.columns)
198    num_data_points = len(df)
199    swce, kwargs_d = get_style_width_color_etc(item_count, num_data_points, **kwargs_d)
200
201    drawn_lines: list[Line2D] = []  # every data line - obstacles for label de-collision
202    annotations: list[tuple[Text, Line2D | None]] = []  # (label, the line it annotates)
203    for i, column in enumerate(df.columns):
204        series = df[column]
205        series = series.dropna() if "dropna" in swce and swce["dropna"][i] else series
206        if series.empty or series.isna().all():
207            print(f"Warning: No data to plot for {column} in line_plot().")
208            continue
209
210        lines = axes.plot(
211            # using matplotlib, as pandas can set xlabel/ylabel
212            series.index,  # x
213            series,  # y
214            ls=swce["style"][i],
215            lw=swce["width"][i],
216            color=swce["color"][i],
217            alpha=swce["alpha"][i],
218            marker=swce["marker"][i],
219            ms=swce["markersize"][i],
220            drawstyle=swce["drawstyle"][i],
221            zorder=swce["zorder"][i],
222            label=(column if "label_series" in swce and swce["label_series"][i] else f"_{column}_"),
223        )
224        drawn_lines.extend(lines)
225
226        if swce["annotate"][i] is None or not swce["annotate"][i]:
227            continue
228
229        color = swce["color"][i] if swce["annotate_color"][i] is True else swce["annotate_color"][i]
230        text = annotate_series(
231            series,
232            axes,
233            color=color,
234            rounding=swce["rounding"][i],
235            fontsize=swce["fontsize"][i],
236            fontname=swce["fontname"][i],
237            rotation=swce["rotation"][i],
238        )
239        if text is not None:
240            annotations.append((text, lines[0] if lines else None))
241
242    # --- register annotations so finalise_plot() can de-collide them
243    if annotations:
244        register_annotations(
245            axes,
246            annotations,
247            drawn_lines,
248            kwargs_d.get("near_end", DEFAULT_NEAR_END),
249        )
250
251    # --- set the labels
252    if saved_pi is not None:
253        set_labels(
254            axes,
255            saved_pi[1],
256            kwargs_d.get("max_ticks", get_setting("max_ticks")),
257            tick_relabel=kwargs_d.get("tick_relabel"),
258        )
259
260    return axes