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 typing import NotRequired, Unpack, cast
  8from pandas import DataFrame, Series, Period, PeriodIndex
  9from matplotlib.pyplot import Axes
 10from numpy import arange, polyfit
 11
 12from mgplot.settings import DataT, get_setting
 13from mgplot.line_plot import line_plot, LineKwargs
 14from mgplot.utilities import check_clean_timeseries
 15from mgplot.keyword_checking import (
 16    validate_kwargs,
 17    report_kwargs,
 18)
 19
 20
 21# --- constants
 22ME = "postcovid_plot"
 23
 24
 25class PostcovidKwargs(LineKwargs):
 26    "Keyword arguments for the post-COVID plot."
 27
 28    start_r: NotRequired[Period]  # start of regression period
 29    end_r: NotRequired[Period]  # end of regression period
 30
 31
 32# --- functions
 33def get_projection(original: Series, to_period: Period) -> Series:
 34    """
 35    Projection based on data from the start of a series
 36    to the to_period (inclusive). Returns projection over the whole
 37    period of the original series.
 38    """
 39
 40    y_regress = original[original.index <= to_period].copy()
 41    x_regress = arange(len(y_regress))
 42    m, b = polyfit(x_regress, y_regress, 1)
 43
 44    x_complete = arange(len(original))
 45    projection = Series((x_complete * m) + b, index=original.index)
 46
 47    return projection
 48
 49
 50def postcovid_plot(data: DataT, **kwargs: Unpack[PostcovidKwargs]) -> Axes:
 51    """
 52    Plots a series with a PeriodIndex.
 53
 54    Arguments
 55    - data - the series to be plotted (note that this function
 56      is designed to work with a single series, not a DataFrame).
 57    - **kwargs - same as for line_plot() and finalise_plot().
 58
 59    Raises:
 60    - TypeError if series is not a pandas Series
 61    - TypeError if series does not have a PeriodIndex
 62    - ValueError if series does not have a D, M or Q frequency
 63    - ValueError if regression start is after regression end
 64    """
 65
 66    # --- check the kwargs
 67    report_kwargs(caller=ME, **kwargs)
 68    validate_kwargs(schema=PostcovidKwargs, caller=ME, **kwargs)
 69
 70    # --- check the data
 71    data = check_clean_timeseries(data, ME)
 72    if not isinstance(data, Series):
 73        raise TypeError("The series argument must be a pandas Series")
 74    series: Series = data
 75    series_index = PeriodIndex(series.index)  # syntactic sugar for type hinting
 76    if series_index.freqstr[:1] not in ("Q", "M", "D"):
 77        raise ValueError("The series index must have a D, M or Q freq")
 78
 79    # rely on line_plot() to validate kwargs
 80    if "plot_from" in kwargs:
 81        print("Warning: the 'plot_from' argument is ignored in postcovid_plot().")
 82        del kwargs["plot_from"]
 83
 84    # --- plot COVID counterfactural
 85    freq = PeriodIndex(series.index).freqstr  # syntactic sugar for type hinting
 86    match freq[0]:
 87        case "Q":
 88            start_regression = Period("2014Q4", freq=freq)
 89            end_regression = Period("2019Q4", freq=freq)
 90        case "M":
 91            start_regression = Period("2015-01", freq=freq)
 92            end_regression = Period("2020-01", freq=freq)
 93        case "D":
 94            start_regression = Period("2015-01-01", freq=freq)
 95            end_regression = Period("2020-01-01", freq=freq)
 96
 97    start_regression = Period(kwargs.pop("start_r", start_regression), freq=freq)
 98    end_regression = Period(kwargs.pop("end_r", end_regression), freq=freq)
 99    if start_regression >= end_regression:
100        raise ValueError("Start period must be before end period")
101
102    # --- combine data and projection
103    recent = series[series.index >= start_regression].copy()
104    recent.name = "Series"
105    projection = get_projection(recent, end_regression)
106    projection.name = "Pre-COVID projection"
107    data_set = DataFrame([projection, recent]).T
108
109    # --- activate plot settings
110    kwargs["width"] = kwargs.pop(
111        "width", (get_setting("line_normal"), get_setting("line_wide"))
112    )  # series line is thicker than projection
113    kwargs["style"] = kwargs.pop("style", ("--", "-"))  # dashed regression line
114    kwargs["label_series"] = kwargs.pop("label_series", True)
115    kwargs["annotate"] = kwargs.pop("annotate", (False, True))  # annotate series only
116    kwargs["color"] = kwargs.pop("color", ("darkblue", "#dd0000"))
117
118    return line_plot(
119        data_set,
120        **cast(LineKwargs, kwargs),
121    )
ME = 'postcovid_plot'
class PostcovidKwargs(mgplot.line_plot.LineKwargs):
26class PostcovidKwargs(LineKwargs):
27    "Keyword arguments for the post-COVID plot."
28
29    start_r: NotRequired[Period]  # start of regression period
30    end_r: NotRequired[Period]  # end of regression period

Keyword arguments for the post-COVID plot.

start_r: NotRequired[pandas._libs.tslibs.period.Period]
end_r: NotRequired[pandas._libs.tslibs.period.Period]
def get_projection( original: pandas.core.series.Series, to_period: pandas._libs.tslibs.period.Period) -> pandas.core.series.Series:
34def get_projection(original: Series, to_period: Period) -> Series:
35    """
36    Projection based on data from the start of a series
37    to the to_period (inclusive). Returns projection over the whole
38    period of the original series.
39    """
40
41    y_regress = original[original.index <= to_period].copy()
42    x_regress = arange(len(y_regress))
43    m, b = polyfit(x_regress, y_regress, 1)
44
45    x_complete = arange(len(original))
46    projection = Series((x_complete * m) + b, index=original.index)
47
48    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: Unpack[PostcovidKwargs]) -> matplotlib.axes._axes.Axes:
 51def postcovid_plot(data: DataT, **kwargs: Unpack[PostcovidKwargs]) -> Axes:
 52    """
 53    Plots a series with a PeriodIndex.
 54
 55    Arguments
 56    - data - the series to be plotted (note that this function
 57      is designed to work with a single series, not a DataFrame).
 58    - **kwargs - same as for line_plot() and finalise_plot().
 59
 60    Raises:
 61    - TypeError if series is not a pandas Series
 62    - TypeError if series does not have a PeriodIndex
 63    - ValueError if series does not have a D, M or Q frequency
 64    - ValueError if regression start is after regression end
 65    """
 66
 67    # --- check the kwargs
 68    report_kwargs(caller=ME, **kwargs)
 69    validate_kwargs(schema=PostcovidKwargs, caller=ME, **kwargs)
 70
 71    # --- check the data
 72    data = check_clean_timeseries(data, ME)
 73    if not isinstance(data, Series):
 74        raise TypeError("The series argument must be a pandas Series")
 75    series: Series = data
 76    series_index = PeriodIndex(series.index)  # syntactic sugar for type hinting
 77    if series_index.freqstr[:1] not in ("Q", "M", "D"):
 78        raise ValueError("The series index must have a D, M or Q freq")
 79
 80    # rely on line_plot() to validate kwargs
 81    if "plot_from" in kwargs:
 82        print("Warning: the 'plot_from' argument is ignored in postcovid_plot().")
 83        del kwargs["plot_from"]
 84
 85    # --- plot COVID counterfactural
 86    freq = PeriodIndex(series.index).freqstr  # syntactic sugar for type hinting
 87    match freq[0]:
 88        case "Q":
 89            start_regression = Period("2014Q4", freq=freq)
 90            end_regression = Period("2019Q4", freq=freq)
 91        case "M":
 92            start_regression = Period("2015-01", freq=freq)
 93            end_regression = Period("2020-01", freq=freq)
 94        case "D":
 95            start_regression = Period("2015-01-01", freq=freq)
 96            end_regression = Period("2020-01-01", freq=freq)
 97
 98    start_regression = Period(kwargs.pop("start_r", start_regression), freq=freq)
 99    end_regression = Period(kwargs.pop("end_r", end_regression), freq=freq)
100    if start_regression >= end_regression:
101        raise ValueError("Start period must be before end period")
102
103    # --- combine data and projection
104    recent = series[series.index >= start_regression].copy()
105    recent.name = "Series"
106    projection = get_projection(recent, end_regression)
107    projection.name = "Pre-COVID projection"
108    data_set = DataFrame([projection, recent]).T
109
110    # --- activate plot settings
111    kwargs["width"] = kwargs.pop(
112        "width", (get_setting("line_normal"), get_setting("line_wide"))
113    )  # series line is thicker than projection
114    kwargs["style"] = kwargs.pop("style", ("--", "-"))  # dashed regression line
115    kwargs["label_series"] = kwargs.pop("label_series", True)
116    kwargs["annotate"] = kwargs.pop("annotate", (False, True))  # annotate series only
117    kwargs["color"] = kwargs.pop("color", ("darkblue", "#dd0000"))
118
119    return line_plot(
120        data_set,
121        **cast(LineKwargs, kwargs),
122    )

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