mgplot 0.1.0__py3-none-any.whl

mgplot/finalisers.py

@@ -0,0 +1,364 @@
+"""
+finalisers.py
+
+Simple convenience functions to finalise and produce plots.
+- bar_plot_finalise()
+- line_plot_finalise()
+- postcovid_plot_finalise()
+- raw_growth_plot_finalise()
+- revision_plot_finalise()
+- run_plot_finalise()
+- seastrend_plot_finalise()
+- series_growth_plot_finalise()
+- summary_plot_finalise()
+
+Note: we keep these functions in a separate module to
+stop circular imports
+
+We also do most of the indicative code testing from this
+module.
+"""
+
+# --- imports
+from pandas import DataFrame, period_range, Period, PeriodIndex, read_csv, Index
+from numpy import random
+
+from mgplot.test import prepare_for_test
+from mgplot.settings import DataT
+from mgplot.multi_plot import plot_then_finalise, multi_column, multi_start
+from mgplot.line_plot import line_plot
+from mgplot.bar_plot import bar_plot
+from mgplot.seastrend_plot import seastrend_plot
+from mgplot.postcovid_plot import postcovid_plot
+from mgplot.revision_plot import revision_plot
+from mgplot.run_plot import run_plot
+from mgplot.growth_plot import series_growth_plot, raw_growth_plot
+from mgplot.summary_plot import summary_plot, ZSCORES, ZSCALED
+
+
+# --- public functions
+def line_plot_finalise(
+    data: DataT,
+    **kwargs,
+) -> None:
+    """
+    A convenience function to call plot_then_finalise(), which
+    wraps calls to line_plot() and finalise_plot().
+    """
+
+    if isinstance(data, DataFrame):
+        if len(data.columns) > 1:
+            # default to displaying a legend
+            kwargs["legend"] = kwargs.get("legend", True)
+        if len(data.columns) > 4:
+            # default to using a style for the lines
+            kwargs["style"] = kwargs.get(
+                "style", ["solid", "dashed", "dashdot", "dotted"]
+            )
+    plot_then_finalise(
+        data,
+        function=line_plot,
+        **kwargs,
+    )
+
+
+def bar_plot_finalise(
+    data: DataT,
+    **kwargs,
+) -> None:
+    """
+    A convenience function to call plot_then_finalise(), which
+    wraps calls to bar_plot() and finalise_plot().
+    """
+
+    plot_then_finalise(
+        data,
+        function=bar_plot,
+        **kwargs,
+    )
+
+
+def seastrend_plot_finalise(
+    data: DataT,
+    **kwargs,
+) -> None:
+    """
+    A convenience function to call seas_trend_plot() and finalise_plot().
+    """
+
+    plot_then_finalise(
+        data,
+        function=seastrend_plot,
+        **kwargs,
+    )
+
+
+def postcovid_plot_finalise(
+    data: DataT,
+    **kwargs,
+) -> None:
+    """
+    A convenience function to call postcovid_plot() and finalise_plot().
+    """
+
+    plot_then_finalise(
+        data,
+        function=postcovid_plot,
+        **kwargs,
+    )
+
+
+def revision_plot_finalise(
+    data: DataT,
+    **kwargs,
+) -> None:
+    """
+    A convenience function to call revision_plot() and finalise_plot().
+    """
+
+    kwargs["legend"] = kwargs.get(
+        "legend", {"loc": "best", "fontsize": "x-small", "ncol": 2}
+    )
+    kwargs["style"] = kwargs.get("style", ["solid", "dashed", "dashdot", "dotted"])
+    plot_then_finalise(
+        data,
+        function=revision_plot,
+        **kwargs,
+    )
+
+
+def run_plot_finalise(
+    data: DataT,
+    **kwargs,
+) -> None:
+    """
+    A convenience function to call run_plot() and finalise_plot().
+    """
+
+    plot_then_finalise(
+        data=data,
+        function=run_plot,
+        **kwargs,
+    )
+
+
+def series_growth_plot_finalise(data: DataT, **kwargs) -> None:
+    """
+    A convenience function to call series_growth_plot() and finalise_plot().
+    Use this when you are providing the series data, for mgplot to calculate
+    the growth series.
+    """
+
+    kwargs["ylabel"] = kwargs.get("ylabel", "Per cent Growth")
+    kwargs["xlabel"] = kwargs.get("xlabel", None)
+    plot_then_finalise(
+        data=data,
+        function=series_growth_plot,
+        **kwargs,
+    )
+
+
+def raw_growth_plot_finalise(data: DataT, **kwargs) -> None:
+    """
+    A convenience function to call series_growth_plot() and finalise_plot().
+    Use this when you are providing the raw growth data. Don't forget to
+    set the ylabel in kwargs.
+    """
+
+    kwargs["ylabel"] = kwargs.get("ylabel", "Growth Units unspecified")
+    kwargs["xlabel"] = kwargs.get("xlabel", None)
+    plot_then_finalise(
+        data=data,
+        function=raw_growth_plot,
+        **kwargs,
+    )
+
+
+def summary_plot_finalise(
+    data: DataT,
+    **kwargs,
+) -> None:
+    """
+    A convenience function to call summary_plot() and finalise_plot().
+    This is more complex than most convienience methods.
+
+    Arguments
+    - data: DataFrame containing the summary data. The index must be a PeriodIndex.
+    - kwargs: additional arguments for the plot, including:
+        - plot_from: int | Period | None  (None means plot from 1995-01-01)
+        - verbose: if True, print the summary data.
+        - middle: proportion of data to highlight (default is 0.8).
+        - plot_type: list of plot types to generate (either "zscores" or "zscaled")
+          defaults to "zscores".
+    """
+
+    # --- sanity checks
+    if not isinstance(data.index, PeriodIndex):
+        raise ValueError("data must have a PeriodIndex")
+
+    # --- standard arguments
+    kwargs["legend"] = kwargs.get(
+        "legend",
+        {
+            # put the legend below the x-axis label
+            "loc": "upper center",
+            "fontsize": "xx-small",
+            "bbox_to_anchor": (0.5, -0.125),
+            "ncol": 4,
+        },
+    )
+    start = kwargs.get("plot_from", None)
+
+    for plot_type in (ZSCORES, ZSCALED):
+        # some sorting of kwargs for plot production
+        kwargs["plot_type"] = plot_type
+        kwargs["title"] = kwargs.get("title", f"Summary at {data.index[-1]}")
+        kwargs["pre_tag"] = plot_type  # necessary because the title is same
+        kwargs["preserve_lims"] = kwargs.get(
+            "preserve_lims", True
+        )  # preserve the x-axis limits
+
+        # get the start date for the plot
+        set_default = start is None
+        if isinstance(start, int):
+            start = data.index[start]
+        if set_default:
+            freq = data.index.freqstr[0]
+            if freq not in ("D", "M", "Q"):
+                raise ValueError(f"Unknown frequency {freq} for data index")
+            start = Period("1995-01-01", freq=data.index.freqstr)
+        kwargs["plot_from"] = start
+
+        if plot_type not in (ZSCORES, ZSCALED):
+            print(f"Unknown plot type {plot_type}, defaulting to {ZSCORES}")
+            plot_type = ZSCORES
+        if plot_type == "zscores":
+            kwargs["xlabel"] = f"Z-scores for prints since {start}"
+            kwargs["x0"] = True
+        else:
+            kwargs["xlabel"] = f"-1 to 1 scaled z-scores since {start}"
+            kwargs.pop("x0", None)
+
+        plot_then_finalise(
+            data,
+            function=summary_plot,
+            **kwargs,
+        )
+
+
+# --- test code
+if __name__ == "__main__":
+    # --- Preparation
+    TEST_DATA_DIR = "./zz-test-data/"
+    prepare_for_test("finalisers")
+
+    # - fake data
+    index = period_range(start="2010Q1", periods=60, freq="Q")
+    test_frame = DataFrame(
+        {
+            "Series 1": [0.1] * len(index),
+            "Series 2": [0.1] * len(index),
+            "Series 3": [1.1] * len(index),
+        },
+        index=index,
+    )
+    test_frame["Series 1"] = test_frame["Series 1"].cumsum() + random.normal(
+        0, 0.1, len(index)
+    )
+    test_frame["Series 2"] = test_frame["Series 2"].cumsum()
+    test_frame["Series 3"] = test_frame["Series 3"].cumprod()
+
+    SKIP = False
+    if not SKIP:
+
+        line_plot_finalise(
+            data=test_frame,
+            title="Test Line Plot",
+            ylabel="Value",
+            xlabel=None,
+        )
+
+        multi_column(
+            data=test_frame,
+            function=line_plot_finalise,
+            title="Test Multi Column Line Plot: ",
+            ylabel="Value",
+            xlabel=None,
+        )
+
+        multi_start(
+            data=test_frame,
+            function=line_plot_finalise,
+            starts=[20, -10, Period("2018Q1")],
+            title="Test Multi Start Line Plot: ",
+            ylabel="Value",
+            xlabel=None,
+        )
+
+        postcovid_plot_finalise(
+            data=test_frame["Series 3"],
+            title="Test Post-COVID Plot",
+            ylabel="Value",
+            xlabel=None,
+        )
+
+        st = test_frame[["Series 1", "Series 2"]].copy()
+        st.columns = Index(["Seasonally Adjusted", "Trend"])
+        multi_start(
+            st,
+            function=seastrend_plot_finalise,
+            starts=[0, Period("2018Q1")],
+            title="Test Multi Start Seas-Trend Plot",
+            ylabel="Value",
+            xlabel=None,
+        )
+
+        # - summary plot test
+        summary_data = read_csv(
+            f"{TEST_DATA_DIR}summary.csv",
+            index_col=0,
+            parse_dates=True,
+        )
+        summary_data.index = PeriodIndex(summary_data.index, freq="M")
+        summary_plot_finalise(
+            data=summary_data,
+            title=f"Summary Plot at {summary_data.index[-1]}",
+            ylabel="Value",
+            xlabel=None,
+        )
+
+        multi_start(
+            data=test_frame["Series 1"],
+            function=series_growth_plot_finalise,
+            starts=[0, -19],
+            title="Test Multi Start Series Growth Plot: ",
+            ylabel="Per cent Growth",
+            xlabel=None,
+        )
+
+        # -- run plot test
+        ocr_data = read_csv(
+            f"{TEST_DATA_DIR}ocr_rba.csv",
+            index_col=0,
+            parse_dates=True,
+        )
+        ocr_data.index = PeriodIndex(ocr_data.index, freq="M")
+        ocr_series = ocr_data[ocr_data.columns[0]]
+        multi_start(
+            data=ocr_series,
+            function=[plot_then_finalise, run_plot],
+            starts=[Period("2020-11", freq="M"), Period("2000-11", freq="M"), 1],
+            title=f"Test Multi Start Run Plot at {ocr_series.index[-1]}",
+            ylabel="Annual Per cent Growth",
+            xlabel=None,
+        )
+
+    data_ = read_csv("./zz-test-data/revisions.csv", index_col=0, parse_dates=True)
+    data_.index = PeriodIndex(data_.index, freq="M")
+    revision_plot_finalise(
+        data=data_,
+        title="Test Revision Plot",
+        ylabel="Units",
+        xlabel=None,
+        rounding=2,
+    )

mgplot/growth_plot.py

@@ -0,0 +1,275 @@
+"""
+growth_plot.py:
+plot period and annual/through-the-year growth rates on the same axes.
+- calc_growth()
+- raw_growth_plot()
+- series_growth_plot()
+"""
+
+# --- imports
+from pandas import Series, DataFrame, Index, Period, PeriodIndex, period_range
+from numpy import nan
+from matplotlib.pyplot import Axes
+import matplotlib.patheffects as pe
+from tabulate import tabulate
+
+from mgplot.test import prepare_for_test
+from mgplot.settings import get_setting, DataT
+from mgplot.date_utils import set_labels
+from mgplot.utilities import annotate_series, check_clean_timeseries
+from mgplot.kw_type_checking import (
+    validate_kwargs,
+    report_kwargs,
+    validate_expected,
+    ExpectedTypeDict,
+)
+
+
+# --- constants
+ANNUAL = "annual"
+PERIODIC = "periodic"
+
+GROWTH_KW_TYPES: ExpectedTypeDict = {
+    "line_width": (float, int),
+    "line_color": str,
+    "line_style": str,
+    "annotate_line": (type(None), int, str),
+    "bar_width": float,
+    "bar_color": str,
+    "annotate_bar": (type(None), int, str),
+    "annotation_rounding": int,
+    "plot_from": (type(None), Period, int),
+    "max_ticks": int,
+}
+validate_expected(GROWTH_KW_TYPES, "growth_plot")
+
+
+# --- functions
+def calc_growth(series: Series) -> DataFrame:
+    """
+    Calculate annual and periodic growth for a pandas Series,
+    where the index is a PeriodIndex.
+
+    Args:
+    -   series: A pandas Series with an appropriate PeriodIndex.
+
+    Returns a two column DataFrame:
+
+    Raises
+    -   TypeError if the series is not a pandas Series.
+    -   TypeError if the series index is not a PeriodIndex.
+    -   ValueError if the series is empty.
+    -   ValueError if the series index does not have a frequency of Q, M, or D.
+    -   ValueError if the series index has duplicates.
+    """
+
+    # --- sanity checks
+    if not isinstance(series, Series):
+        raise TypeError("The series argument must be a pandas Series")
+    if not isinstance(series.index, PeriodIndex):
+        raise TypeError("The series index must be a pandas PeriodIndex")
+    if series.empty:
+        raise ValueError("The series argument must not be empty")
+    if series.index.freqstr[0] not in ("Q", "M", "D"):
+        raise ValueError("The series index must have a frequency of Q, M, or D")
+    if series.index.has_duplicates:
+        raise ValueError("The series index must not have duplicate values")
+
+    # --- ensure the index is complete and the date is sorted
+    complete = period_range(start=series.index.min(), end=series.index.max())
+    series = series.reindex(complete, fill_value=nan)
+    series = series.sort_index(ascending=True)
+
+    # --- calculate annual and periodic growth
+    ppy = {"Q": 4, "M": 12, "D": 365}[PeriodIndex(series.index).freqstr[:1]]
+    annual = series.pct_change(periods=ppy) * 100
+    periodic = series.pct_change(periods=1) * 100
+    periodic_name = {4: "Quarterly", 12: "Monthly", 365: "Daily"}[ppy] + " Growth"
+    return DataFrame(
+        {
+            "Annual Growth": annual,
+            periodic_name: periodic,
+        }
+    )
+
+
+def _annotations(
+    annual: Series,
+    periodic: Series,
+    axes: Axes,
+    **kwargs,
+) -> None:
+    """Apply annotations the annual and periodic growth series."""
+
+    annotate_line = kwargs.get("annotate_line", "small")
+    if annotate_line is not None:
+        annotate_series(
+            annual,
+            axes,
+            rounding=kwargs.get("annotation_rounding", True),
+            fontsize=annotate_line,
+            color=kwargs.get("line_color", "darkblue"),
+        )
+
+    annotate_bar = kwargs.get("annotate_bar", "small")
+    max_annotations = 30
+    if annotate_bar is not None and len(periodic) < max_annotations:
+        annotation_rounding = kwargs.get("annotation_rounding", 1)
+        annotate_style = {
+            "fontsize": annotate_bar,
+            "fontname": "Helvetica",
+        }
+        adjustment = (periodic.max() - periodic.min()) * 0.005
+        for i, value in enumerate(periodic):
+            va = "bottom" if value >= 0 else "top"
+            text = axes.text(
+                periodic.index[i],
+                adjustment if value >= 0 else -adjustment,
+                f"{value:.{annotation_rounding}f}",
+                ha="center",
+                va=va,
+                **annotate_style,
+                fontdict=None,
+                color="white",
+            )
+            text.set_path_effects(
+                [
+                    pe.withStroke(
+                        linewidth=2, foreground=kwargs.get("bar_color", "indianred")
+                    )
+                ]
+            )
+
+
+def raw_growth_plot(
+    data: DataT,
+    **kwargs,
+) -> Axes:
+    """
+    Plot annual (as a line) and periodic (as bars) growth on the
+    same axes.
+
+    Args:
+    -   data: A pandas DataFrame with two columns:
+    -   kwargs:
+        -   line_width: The width of the line (default is 2).
+        -   line_color: The color of the line (default is "darkblue").
+        -   line_style: The style of the line (default is "-").
+        -   annotate_line: None | int | str - fontsize to annotate the line
+            (default is "small", which means the line is annotated with
+            small text).
+        -   bar_width: The width of the bars (default is 0.8).
+        -   bar_color: The color of the bars (default is "indianred").
+        -   annotate_bar: None | int | str - fontsize to annotate the bars
+            (default is "small", which means the bars are annotated with
+            small text).
+        -   annotation_rounding: The number of decimal places to round the
+            annotations to (default is 1).
+        -   plot_from: None | Period | int -- if:
+            -   None: the entire series is plotted
+            -   Period: the plot starts from this period
+            -   int: the plot starts from this +/- index position
+        -   max_ticks: The maximum number of ticks to show on the x-axis
+            (default is 10).
+
+    Returns:
+    -   axes: The matplotlib Axes object.
+
+    Raises:
+    -   TypeError if the annual and periodic arguments are not pandas Series.
+    -   TypeError if the annual index is not a PeriodIndex.
+    -   ValueError if the annual and periodic series do not have the same index.
+    """
+
+    # --- sanity checks
+    report_kwargs(called_from="raw_growth_plot", **kwargs)
+    validate_kwargs(GROWTH_KW_TYPES, "raw_growth_plot", **kwargs)
+    data = check_clean_timeseries(data)
+    if len(data.columns) != 2:
+        raise TypeError("The data argument must be a pandas DataFrame with two columns")
+
+    # --- get the series of interest ...
+    annual = data[data.columns[0]]
+    periodic = data[data.columns[1]]
+
+    # --- plot
+    plot_from: None | Period | int = kwargs.get("plot_from", None)
+    if plot_from is not None:
+        if isinstance(plot_from, int):
+            plot_from = annual.index[plot_from]
+        annual = annual[annual.index >= plot_from]
+        periodic = periodic[periodic.index >= plot_from]
+
+    save_index = PeriodIndex(annual.index).copy()
+    annual.index = Index(range(len(annual)))
+    annual.name = "Annual Growth"
+    periodic.index = annual.index
+    periodic.name = {"M": "Monthly", "Q": "Quarterly", "D": "Daily"}[
+        PeriodIndex(save_index).freqstr[:1]
+    ] + " Growth"
+    axes = periodic.plot.bar(
+        color=kwargs.get("bar_color", "indianred"),
+        width=kwargs.get("bar_width}", 0.8),
+    )
+    thin_threshold = 180
+    annual.plot(
+        ax=axes,
+        color=kwargs.get("line_color", "darkblue"),
+        lw=kwargs.get(
+            "line_width",
+            (
+                get_setting("line_normal")
+                if len(annual) >= thin_threshold
+                else get_setting("line_wide")
+            ),
+        ),
+        linestyle=kwargs.get("line_style", "-"),
+    )
+    _annotations(annual, periodic, axes, **kwargs)
+    axes.set_ylabel("Per cent Growth")
+
+    # --- fix the x-axis labels
+    set_labels(axes, save_index, kwargs.get("max_ticks", 10))
+
+    # --- and done ...
+    return axes
+
+
+def series_growth_plot(
+    data: DataT,
+    **kwargs,
+) -> Axes:
+    """
+    Plot annual and periodic growth from a pandas Series,
+    and finalise the plot.
+
+    Args:
+    -   data: A pandas Series with an appropriate PeriodIndex.
+    -   kwargs:
+        -   takes the same kwargs as for growth_plot()
+    """
+
+    # --- sanity checks
+    report_kwargs(called_from="series_growth_plot", **kwargs)
+    data = check_clean_timeseries(data)
+    # we will validate kwargs in raw_growth_plot()
+    if not isinstance(data, Series):
+        raise TypeError(
+            "The data argument to series_growth_plot() must be a pandas Series"
+        )
+
+    # --- calculate growth and plot
+    growth = calc_growth(data)
+    ax = raw_growth_plot(growth, **kwargs)
+    return ax
+
+
+# --- test code
+if __name__ == "__main__":
+    print("Testing")
+    prepare_for_test("growth_plot")
+    series_ = Series([1, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0])
+    series_.index = period_range("2020Q1", periods=len(series_), freq="Q")
+    growth_ = calc_growth(series_)
+    text_ = tabulate(growth_, headers="keys", tablefmt="pipe")  # type: ignore[arg-type]
+    print(text_)