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