mgplot.postcovid_plot

  1"""Plot the linear pre-COVID trajectory against the current data."""
  2
  3from typing import Literal, NotRequired, Unpack
  4
  5from matplotlib.axes import Axes
  6from numpy import array, polyfit
  7from pandas import DataFrame, Period, PeriodIndex, Series, period_range
  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"
 19MIN_REGRESSION_POINTS = 10  # minimum number of points for a useful linear regression
 20
 21# Default regression periods by frequency
 22DEFAULT_PERIODS = {
 23    "Q": {"start": "2014Q4", "end": "2019Q4"},
 24    "M": {"start": "2015-01", "end": "2020-01"},
 25    "D": {"start": "2015-01-01", "end": "2020-01-01"},
 26}
 27
 28
 29class PostcovidKwargs(LineKwargs):
 30    """Keyword arguments for the post-COVID plot."""
 31
 32    start_r: NotRequired[Period]  # start of regression period
 33    end_r: NotRequired[Period]  # end of regression period
 34
 35
 36# --- functions
 37def get_projection(source: Series, to_period: Period) -> Series:
 38    """Create a linear projection based on pre-COVID data.
 39
 40    Args:
 41        source: Series - the original series with a PeriodIndex
 42            Assume the index is a PeriodIndex, that is unique and monotonic increasing.
 43            Assume there may be gaps in the source series (either missing or NaNs)
 44            And that it starts from when the regression should start.
 45        to_period: Period - the period to which the projection should extend.
 46
 47    Returns:
 48        Series: A pandas Series with linear projection values using the same index as original.
 49            Returns an empty Series if it fails to create a projection.
 50
 51    Raises:
 52        ValueError: If to_period is not within the original series index range.
 53
 54    """
 55    # --- initial validation
 56    if not isinstance(source.index, PeriodIndex):
 57        raise TypeError("Source index must be a PeriodIndex")
 58    if source.empty or not source.index.is_monotonic_increasing or not source.index.is_unique:
 59        print("Source series must be non-empty, uniquely indexed, and a monotonic increasing index.")
 60        return Series(dtype=float)  # return empty series if validation fails
 61
 62    # --- Drop any missing data and establish the input data for regression
 63    source_no_nan = source.dropna()
 64    input_series = source_no_nan[source_no_nan.index <= to_period]
 65
 66    # --- further validation
 67    if input_series.empty or len(input_series) < MIN_REGRESSION_POINTS:
 68        print("Insufficient data points for regression.")
 69        return Series(dtype=float)  # return empty series if no data for regression
 70
 71    # --- Establish the simple linear regression model
 72    input_index = input_series.index
 73    x_cause = array([p.ordinal for p in input_index if p <= to_period])
 74    y_effect = input_series.to_numpy()
 75    slope, intercept = polyfit(x_cause, y_effect, 1)
 76
 77    # --- use the regression model to create an out-of-sample projection
 78    x_complete = array([p.ordinal for p in source.index])
 79    projection = Series((x_complete * slope) + intercept, index=source.index)
 80
 81    # --- ensure the projection covers any date gaps in the PeriodIndex
 82    source_index = source.index
 83    return projection.reindex(period_range(start=source_index[0], end=source_index[-1])).interpolate(
 84        method="linear"
 85    )
 86
 87
 88def regression_period(data: Series, **kwargs: Unpack[PostcovidKwargs]) -> tuple[Period, Period, bool]:
 89    """Establish the regression period.
 90
 91    Args:
 92        data: Series - the original time series data.
 93        **kwargs: Additional keyword arguments.
 94
 95    Returns:
 96        A tuple containing the start and end periods for regression,
 97        and a boolean indicating if the period is robust.
 98
 99    Raises:
100        TypeError: If the series index is not a PeriodIndex.
101        ValueError: If the series index does not have a D, M, or Q frequency
102
103    """
104    # --- check that the series index is a PeriodIndex with a valid frequency
105    if not isinstance(data.index, PeriodIndex):
106        raise TypeError("The series index must be a PeriodIndex")
107    freq_str = data.index.freqstr
108    freq_key = freq_str[0]
109    if not freq_str or freq_key not in ("Q", "M", "D"):
110        raise ValueError("The series index must have a D, M or Q frequency")
111
112    # --- set the default regression period, use user provided periods if specified
113    default_periods = DEFAULT_PERIODS[freq_key]
114    start_regression = Period(default_periods["start"], freq=freq_str)
115    end_regression = Period(default_periods["end"], freq=freq_str)
116
117    user_start = kwargs.pop("start_r", None)
118    user_end = kwargs.pop("end_r", None)
119    start_r = Period(user_start, freq=freq_str) if user_start else start_regression
120    end_r = Period(user_end, freq=freq_str) if user_end else end_regression
121
122    # --- Validate the regression period
123    robust = True
124    if start_r >= end_r:
125        print(f"Invalid regression period: {start_r=}, {end_r=}")
126        robust = False
127
128    return start_r, end_r, robust
129
130
131def postcovid_plot(data: DataT, **kwargs: Unpack[PostcovidKwargs]) -> Axes:
132    """Plot a series with a PeriodIndex, including a post-COVID projection.
133
134    Args:
135        data: Series - the series to be plotted.
136        kwargs: PostcovidKwargs - plotting arguments.
137
138    Raises:
139        TypeError if series is not a pandas Series
140        TypeError if series does not have a PeriodIndex
141        ValueError if series does not have a D, M or Q frequency
142        ValueError if regression start is after regression end
143
144    """
145
146    # --- failure
147    def failure() -> Axes:
148        print("postcovid_plot(): plotting the raw data only.")
149        remove: list[Literal["plot_from", "start_r", "end_r"]] = ["plot_from", "start_r", "end_r"]
150        for key in remove:
151            kwargs.pop(key, None)
152        line_kwargs: LineKwargs = kwargs  # start_r, end_r and plot_from removed above
153        return line_plot(data, **line_kwargs)
154
155    # --- check the kwargs
156    report_kwargs(caller=ME, **kwargs)
157    validate_kwargs(schema=PostcovidKwargs, caller=ME, **kwargs)
158
159    # --- check the data
160    data = check_clean_timeseries(data, ME)
161    if not isinstance(data, Series):
162        raise TypeError("The series argument must be a pandas Series")
163
164    # --- rely on line_plot() to validate kwargs, but remove any that are not relevant
165    if "plot_from" in kwargs:
166        print("Warning: the 'plot_from' argument is ignored in postcovid_plot().")
167        kwargs.pop("plot_from", None)
168
169    # --- set the regression period
170    start_r, end_r, robust = regression_period(data, **kwargs)
171    kwargs.pop("start_r", None)  # remove from kwargs to avoid confusion
172    kwargs.pop("end_r", None)  # remove from kwargs to avoid confusion
173    if not robust:
174        return failure()
175
176    # --- combine data and projection
177    if start_r < data.dropna().index.min():
178        print(f"Caution: Regression start period pre-dates the series index: {start_r=}")
179    recent_data = data[data.index >= start_r].copy()
180    recent_data.name = "Series"
181    projection_data = get_projection(recent_data, end_r)
182    if projection_data.empty:
183        return failure()
184    projection_data.name = "Pre-COVID projection"
185
186    # --- Create DataFrame with proper column alignment
187    combined_data = DataFrame(
188        {
189            projection_data.name: projection_data,
190            recent_data.name: recent_data,
191        }
192    )
193
194    # --- activate plot settings
195    kwargs["width"] = kwargs.pop(
196        "width",
197        (get_setting("line_normal"), get_setting("line_wide")),
198    )  # series line is thicker than projection
199    kwargs["style"] = kwargs.pop("style", ("--", "-"))  # dashed regression line
200    kwargs["label_series"] = kwargs.pop("label_series", True)
201    kwargs["annotate"] = kwargs.pop("annotate", (False, True))  # annotate series only
202    kwargs["color"] = kwargs.pop("color", ("darkblue", "#dd0000"))
203    kwargs["dropna"] = kwargs.pop("dropna", False)  # drop NaN values
204
205    line_kwargs: LineKwargs = kwargs  # start_r, end_r and plot_from removed above
206    return line_plot(combined_data, **line_kwargs)
207
208
209if __name__ == "__main__":
210
211    def test_make_projection() -> None:
212        """Test the get_projection function."""
213        n = 30
214        periods = period_range(start="2015-Q1", periods=n, freq="Q")
215        series = Series(
216            [i + (i % 3) for i in range(n)],  # simple increasing series with some noise
217            index=periods,
218        )
219        proj = get_projection(series, Period("2019-Q4", freq="Q"))
220        print(
221            DataFrame(
222                {
223                    "Input": series,
224                    "Projection": proj,
225                }
226            )
227        )
228
229    test_make_projection()