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'
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.
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.
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