mgplot.postcovid_plot

covid_recovery_plot.py Plot the pre-COVID trajectory against the current trend.

  1"""
  2covid_recovery_plot.py
  3Plot the pre-COVID trajectory against the current trend.
  4"""
  5
  6# --- imports
  7from pandas import DataFrame, Series, Period, PeriodIndex
  8from matplotlib.pyplot import Axes
  9from numpy import arange, polyfit
 10
 11from mgplot.settings import DataT, get_setting
 12from mgplot.line_plot import line_plot, LINE_KW_TYPES
 13from mgplot.utilities import check_clean_timeseries
 14from mgplot.kw_type_checking import (
 15    ExpectedTypeDict,
 16    validate_kwargs,
 17    validate_expected,
 18    report_kwargs,
 19)
 20
 21
 22# --- constants
 23START_R = "start_r"
 24END_R = "end_r"
 25WIDTH = "width"
 26STYLE = "style"
 27
 28POSTCOVID_KW_TYPES: ExpectedTypeDict = {
 29    START_R: Period,
 30    END_R: Period,
 31} | LINE_KW_TYPES
 32validate_expected(POSTCOVID_KW_TYPES, "postcovid_plot")
 33
 34
 35# --- functions
 36def get_projection(original: Series, to_period: Period) -> Series:
 37    """
 38    Projection based on data from the start of a series
 39    to the to_period (inclusive). Returns projection over the whole
 40    period of the original series.
 41    """
 42
 43    y_regress = original[original.index <= to_period].copy()
 44    x_regress = arange(len(y_regress))
 45    m, b = polyfit(x_regress, y_regress, 1)
 46
 47    x_complete = arange(len(original))
 48    projection = Series((x_complete * m) + b, index=original.index)
 49
 50    return projection
 51
 52
 53def postcovid_plot(data: DataT, **kwargs) -> Axes:
 54    """
 55    Plots a series with a PeriodIndex.
 56
 57    Arguments
 58    - data - the series to be plotted (note that this function
 59      is designed to work with a single series, not a DataFrame).
 60    - **kwargs - same as for line_plot() and finalise_plot().
 61
 62    Raises:
 63    - TypeError if series is not a pandas Series
 64    - TypeError if series does not have a PeriodIndex
 65    - ValueError if series does not have a D, M or Q frequency
 66    - ValueError if regression start is after regression end
 67    """
 68
 69    # --- check the kwargs
 70    me = "postcovid_plot"
 71    report_kwargs(called_from=me, **kwargs)
 72    validate_kwargs(POSTCOVID_KW_TYPES, me, **kwargs)
 73
 74    # --- check the data
 75    data = check_clean_timeseries(data, me)
 76    if not isinstance(data, Series):
 77        raise TypeError("The series argument must be a pandas Series")
 78    series: Series = data
 79    series_index = PeriodIndex(series.index)  # syntactic sugar for type hinting
 80    if series_index.freqstr[:1] not in ("Q", "M", "D"):
 81        raise ValueError("The series index must have a D, M or Q freq")
 82    # rely on line_plot() to validate kwargs
 83    if "plot_from" in kwargs:
 84        print("Warning: the 'plot_from' argument is ignored in postcovid_plot().")
 85        del kwargs["plot_from"]
 86
 87    # --- plot COVID counterfactural
 88    freq = PeriodIndex(series.index).freqstr  # syntactic sugar for type hinting
 89    match freq[0]:
 90        case "Q":
 91            start_regression = Period("2014Q4", freq=freq)
 92            end_regression = Period("2019Q4", freq=freq)
 93        case "M":
 94            start_regression = Period("2015-01", freq=freq)
 95            end_regression = Period("2020-01", freq=freq)
 96        case "D":
 97            start_regression = Period("2015-01-01", freq=freq)
 98            end_regression = Period("2020-01-01", freq=freq)
 99
100    start_regression = Period(kwargs.pop("start_r", start_regression), freq=freq)
101    end_regression = Period(kwargs.pop("end_r", end_regression), freq=freq)
102    if start_regression >= end_regression:
103        raise ValueError("Start period must be before end period")
104
105    # --- combine data and projection
106    recent = series[series.index >= start_regression].copy()
107    recent.name = "Series"
108    projection = get_projection(recent, end_regression)
109    projection.name = "Pre-COVID projection"
110    data_set = DataFrame([projection, recent]).T
111
112    # --- activate plot settings
113    kwargs[WIDTH] = kwargs.pop(
114        WIDTH, (get_setting("line_normal"), get_setting("line_wide"))
115    )  # series line is thicker than projection
116    kwargs[STYLE] = kwargs.pop(STYLE, ("--", "-"))  # dashed regression line
117    kwargs["legend"] = kwargs.pop("legend", True)  # show legend by default
118    kwargs["annotate"] = kwargs.pop("annotate", (False, True))  # annotate series only
119    kwargs["color"] = kwargs.pop("color", ("darkblue", "#dd0000"))
120
121    return line_plot(
122        data_set,
123        **kwargs,
124    )
START_R = 'start_r'
END_R = 'end_r'
WIDTH = 'width'
STYLE = 'style'
POSTCOVID_KW_TYPES: mgplot.kw_type_checking.ExpectedTypeDict = {'start_r': <class 'pandas._libs.tslibs.period.Period'>, 'end_r': <class 'pandas._libs.tslibs.period.Period'>, 'ax': (<class 'matplotlib.axes._axes.Axes'>, <class 'NoneType'>), 'style': (<class 'str'>, <class 'collections.abc.Sequence'>, (<class 'str'>,)), 'width': (<class 'float'>, <class 'int'>, <class 'collections.abc.Sequence'>, (<class 'float'>, <class 'int'>)), 'color': (<class 'str'>, <class 'collections.abc.Sequence'>, (<class 'str'>,)), 'alpha': (<class 'float'>, <class 'collections.abc.Sequence'>, (<class 'float'>,)), 'drawstyle': (<class 'str'>, <class 'collections.abc.Sequence'>, (<class 'str'>,), <class 'NoneType'>), 'marker': (<class 'str'>, <class 'collections.abc.Sequence'>, (<class 'str'>,), <class 'NoneType'>), 'markersize': (<class 'float'>, <class 'collections.abc.Sequence'>, (<class 'float'>,), <class 'int'>, <class 'NoneType'>), 'dropna': (<class 'bool'>, <class 'collections.abc.Sequence'>, (<class 'bool'>,)), 'annotate': (<class 'bool'>, <class 'collections.abc.Sequence'>, (<class 'bool'>,)), 'rounding': (<class 'collections.abc.Sequence'>, (<class 'bool'>, <class 'int'>), <class 'int'>, <class 'bool'>, <class 'NoneType'>), 'fontsize': (<class 'collections.abc.Sequence'>, (<class 'str'>, <class 'int'>), <class 'str'>, <class 'int'>, <class 'NoneType'>), 'plot_from': (<class 'int'>, <class 'pandas._libs.tslibs.period.Period'>, <class 'NoneType'>), 'legend': (<class 'dict'>, (<class 'str'>, <class 'object'>), <class 'bool'>, <class 'NoneType'>)}
def get_projection( original: pandas.core.series.Series, to_period: pandas._libs.tslibs.period.Period) -> pandas.core.series.Series:
37def get_projection(original: Series, to_period: Period) -> Series:
38    """
39    Projection based on data from the start of a series
40    to the to_period (inclusive). Returns projection over the whole
41    period of the original series.
42    """
43
44    y_regress = original[original.index <= to_period].copy()
45    x_regress = arange(len(y_regress))
46    m, b = polyfit(x_regress, y_regress, 1)
47
48    x_complete = arange(len(original))
49    projection = Series((x_complete * m) + b, index=original.index)
50
51    return projection

Projection based on data from the start of a series to the to_period (inclusive). Returns projection over the whole period of the original series.

def postcovid_plot(data: ~DataT, **kwargs) -> matplotlib.axes._axes.Axes:
 54def postcovid_plot(data: DataT, **kwargs) -> Axes:
 55    """
 56    Plots a series with a PeriodIndex.
 57
 58    Arguments
 59    - data - the series to be plotted (note that this function
 60      is designed to work with a single series, not a DataFrame).
 61    - **kwargs - same as for line_plot() and finalise_plot().
 62
 63    Raises:
 64    - TypeError if series is not a pandas Series
 65    - TypeError if series does not have a PeriodIndex
 66    - ValueError if series does not have a D, M or Q frequency
 67    - ValueError if regression start is after regression end
 68    """
 69
 70    # --- check the kwargs
 71    me = "postcovid_plot"
 72    report_kwargs(called_from=me, **kwargs)
 73    validate_kwargs(POSTCOVID_KW_TYPES, me, **kwargs)
 74
 75    # --- check the data
 76    data = check_clean_timeseries(data, me)
 77    if not isinstance(data, Series):
 78        raise TypeError("The series argument must be a pandas Series")
 79    series: Series = data
 80    series_index = PeriodIndex(series.index)  # syntactic sugar for type hinting
 81    if series_index.freqstr[:1] not in ("Q", "M", "D"):
 82        raise ValueError("The series index must have a D, M or Q freq")
 83    # rely on line_plot() to validate kwargs
 84    if "plot_from" in kwargs:
 85        print("Warning: the 'plot_from' argument is ignored in postcovid_plot().")
 86        del kwargs["plot_from"]
 87
 88    # --- plot COVID counterfactural
 89    freq = PeriodIndex(series.index).freqstr  # syntactic sugar for type hinting
 90    match freq[0]:
 91        case "Q":
 92            start_regression = Period("2014Q4", freq=freq)
 93            end_regression = Period("2019Q4", freq=freq)
 94        case "M":
 95            start_regression = Period("2015-01", freq=freq)
 96            end_regression = Period("2020-01", freq=freq)
 97        case "D":
 98            start_regression = Period("2015-01-01", freq=freq)
 99            end_regression = Period("2020-01-01", freq=freq)
100
101    start_regression = Period(kwargs.pop("start_r", start_regression), freq=freq)
102    end_regression = Period(kwargs.pop("end_r", end_regression), freq=freq)
103    if start_regression >= end_regression:
104        raise ValueError("Start period must be before end period")
105
106    # --- combine data and projection
107    recent = series[series.index >= start_regression].copy()
108    recent.name = "Series"
109    projection = get_projection(recent, end_regression)
110    projection.name = "Pre-COVID projection"
111    data_set = DataFrame([projection, recent]).T
112
113    # --- activate plot settings
114    kwargs[WIDTH] = kwargs.pop(
115        WIDTH, (get_setting("line_normal"), get_setting("line_wide"))
116    )  # series line is thicker than projection
117    kwargs[STYLE] = kwargs.pop(STYLE, ("--", "-"))  # dashed regression line
118    kwargs["legend"] = kwargs.pop("legend", True)  # show legend by default
119    kwargs["annotate"] = kwargs.pop("annotate", (False, True))  # annotate series only
120    kwargs["color"] = kwargs.pop("color", ("darkblue", "#dd0000"))
121
122    return line_plot(
123        data_set,
124        **kwargs,
125    )

Plots a series with a PeriodIndex.

Arguments

  • data - the series to be plotted (note that this function is designed to work with a single series, not a DataFrame).
  • **kwargs - same as for line_plot() and finalise_plot().

Raises:

  • TypeError if series is not a pandas Series
  • TypeError if series does not have a PeriodIndex
  • ValueError if series does not have a D, M or Q frequency
  • ValueError if regression start is after regression end