mgplot.postcovid_plot

  1"""Plot the pre-COVID trajectory against the current trend."""
  2
  3from typing import NotRequired, Unpack, cast
  4
  5from matplotlib.axes import Axes
  6from numpy import arange, polyfit
  7from pandas import DataFrame, Period, PeriodIndex, Series
  8
  9from mgplot.keyword_checking import (
 10    report_kwargs,
 11    validate_kwargs,
 12)
 13from mgplot.line_plot import LineKwargs, line_plot
 14from mgplot.settings import DataT, get_setting
 15from mgplot.utilities import check_clean_timeseries
 16
 17# --- constants
 18ME = "postcovid_plot"
 19
 20# Default regression periods by frequency
 21DEFAULT_PERIODS = {
 22    "Q": {"start": "2014Q4", "end": "2019Q4"},
 23    "M": {"start": "2015-01", "end": "2020-01"},
 24    "D": {"start": "2015-01-01", "end": "2020-01-01"},
 25}
 26
 27
 28class PostcovidKwargs(LineKwargs):
 29    """Keyword arguments for the post-COVID plot."""
 30
 31    start_r: NotRequired[Period]  # start of regression period
 32    end_r: NotRequired[Period]  # end of regression period
 33
 34
 35# --- functions
 36def get_projection(original: Series, to_period: Period) -> Series:
 37    """Create a linear projection based on pre-COVID data.
 38
 39    Assumes the start of the data has been trimmed to the period before COVID.
 40
 41    Args:
 42        original: Series - the original series with a PeriodIndex.
 43        to_period: Period - the period to which the projection should extend.
 44
 45    Returns:
 46        Series: A pandas Series with linear projection values using the same index as original.
 47
 48    Raises:
 49        ValueError: If to_period is not within the original series index range.
 50
 51    """
 52    if to_period not in original.index:
 53        raise ValueError(f"Regression end period {to_period} not found in series index")
 54    y_regress = original[original.index <= to_period].copy()
 55    x_regress = arange(len(y_regress))
 56    m, b = polyfit(x_regress, y_regress, 1)
 57
 58    x_complete = arange(len(original))
 59    return Series((x_complete * m) + b, index=original.index)
 60
 61
 62def postcovid_plot(data: DataT, **kwargs: Unpack[PostcovidKwargs]) -> Axes:
 63    """Plot a series with a PeriodIndex, including a post-COVID projection.
 64
 65    Args:
 66        data: Series - the series to be plotted.
 67        kwargs: PostcovidKwargs - plotting arguments.
 68
 69    Raises:
 70        TypeError if series is not a pandas Series
 71        TypeError if series does not have a PeriodIndex
 72        ValueError if series does not have a D, M or Q frequency
 73        ValueError if regression start is after regression end
 74
 75    """
 76    # --- check the kwargs
 77    report_kwargs(caller=ME, **kwargs)
 78    validate_kwargs(schema=PostcovidKwargs, caller=ME, **kwargs)
 79
 80    # --- check the data
 81    data = check_clean_timeseries(data, ME)
 82    if not isinstance(data, Series):
 83        raise TypeError("The series argument must be a pandas Series")
 84
 85    series_index = PeriodIndex(data.index)
 86    freq_str = series_index.freqstr
 87    if not freq_str or freq_str[0] not in ("Q", "M", "D"):
 88        raise ValueError("The series index must have a D, M or Q frequency")
 89
 90    freq_key = freq_str[0]
 91
 92    # rely on line_plot() to validate kwargs
 93    if "plot_from" in kwargs:
 94        print("Warning: the 'plot_from' argument is ignored in postcovid_plot().")
 95        del kwargs["plot_from"]
 96
 97    # --- plot COVID counterfactual
 98    default_periods = DEFAULT_PERIODS[freq_key]
 99    start_regression = Period(default_periods["start"], freq=freq_str)
100    end_regression = Period(default_periods["end"], freq=freq_str)
101
102    # Override defaults with user-provided periods if specified
103    user_start = kwargs.pop("start_r", None)
104    user_end = kwargs.pop("end_r", None)
105
106    if user_start is not None:
107        start_regression = Period(user_start, freq=freq_str)
108    if user_end is not None:
109        end_regression = Period(user_end, freq=freq_str)
110
111    # Validate regression period
112    if start_regression >= end_regression:
113        raise ValueError("Start period must be before end period")
114
115    if start_regression not in data.index:
116        raise ValueError(f"Regression start period {start_regression} not found in series")
117    if end_regression not in data.index:
118        raise ValueError(f"Regression end period {end_regression} not found in series")
119
120    # --- combine data and projection
121    recent_data = data[data.index >= start_regression].copy()
122    recent_data.name = "Series"
123    projection_data = get_projection(recent_data, end_regression)
124    projection_data.name = "Pre-COVID projection"
125
126    # Create DataFrame with proper column alignment
127    combined_data = DataFrame(
128        {
129            projection_data.name: projection_data,
130            recent_data.name: recent_data,
131        }
132    )
133
134    # --- activate plot settings
135    kwargs["width"] = kwargs.pop(
136        "width",
137        (get_setting("line_normal"), get_setting("line_wide")),
138    )  # series line is thicker than projection
139    kwargs["style"] = kwargs.pop("style", ("--", "-"))  # dashed regression line
140    kwargs["label_series"] = kwargs.pop("label_series", True)
141    kwargs["annotate"] = kwargs.pop("annotate", (False, True))  # annotate series only
142    kwargs["color"] = kwargs.pop("color", ("darkblue", "#dd0000"))
143
144    return line_plot(
145        combined_data,
146        **cast("LineKwargs", kwargs),
147    )