Coverage for src / ts_stat_tests / linearity / tests.py: 100%

1# ============================================================================ # 

2# # 

3# Title: Linearity Tests # 

4# Purpose: Convenience functions for linearity algorithms. # 

5# # 

6# ============================================================================ # 

7 

8 

9# ---------------------------------------------------------------------------- # 

10# # 

11# Overview #### 

12# # 

13# ---------------------------------------------------------------------------- # 

14 

15 

16# ---------------------------------------------------------------------------- # 

17# Description #### 

18# ---------------------------------------------------------------------------- # 

19 

20 

21""" 

22!!! note "Summary" 

23 This module contains convenience functions and tests for linearity measures, allowing for easy access to different linearity algorithms. 

24""" 

25 

26 

27# ---------------------------------------------------------------------------- # 

28# # 

29# Setup #### 

30# # 

31# ---------------------------------------------------------------------------- # 

32 

33 

34# ---------------------------------------------------------------------------- # 

35# Imports #### 

36# ---------------------------------------------------------------------------- # 

37 

38 

39# ## Python StdLib Imports ---- 

40from typing import Callable, Union 

41 

42# ## Python Third Party Imports ---- 

43import numpy as np 

44from numpy.typing import ArrayLike 

45from statsmodels.regression.linear_model import ( 

46 RegressionResults, 

47 RegressionResultsWrapper, 

48) 

49from statsmodels.stats.contrast import ContrastResults 

50from typeguard import typechecked 

51 

52# ## Local First Party Imports ---- 

53from ts_stat_tests.linearity.algorithms import ( 

54 hc as _hc, 

55 lm as _lm, 

56 rb as _rb, 

57 rr as _rr, 

58) 

59from ts_stat_tests.utils.errors import generate_error_message 

60 

61 

62# ---------------------------------------------------------------------------- # 

63# Exports #### 

64# ---------------------------------------------------------------------------- # 

65 

66 

67__all__: list[str] = ["linearity", "is_linear"] 

68 

69 

70# ---------------------------------------------------------------------------- # 

71# # 

72# Tests #### 

73# # 

74# ---------------------------------------------------------------------------- # 

75 

76 

77@typechecked 

78def linearity( 

79 res: Union[RegressionResults, RegressionResultsWrapper], 

80 algorithm: str = "rr", 

81 **kwargs: Union[float, int, str, bool, ArrayLike, None], 

82) -> Union[tuple[float, ...], ContrastResults]: 

83 """ 

84 !!! note "Summary" 

85 Perform a linearity test on a fitted regression model. 

86 

87 ???+ abstract "Details" 

88 This function is a convenience wrapper around four underlying algorithms:<br> 

89 - [`hc()`][ts_stat_tests.linearity.algorithms.hc]<br> 

90 - [`lm()`][ts_stat_tests.linearity.algorithms.lm]<br> 

91 - [`rb()`][ts_stat_tests.linearity.algorithms.rb]<br> 

92 - [`rr()`][ts_stat_tests.linearity.algorithms.rr] 

93 

94 Params: 

95 res (Union[RegressionResults, RegressionResultsWrapper]): 

96 The fitted regression model to be checked. 

97 algorithm (str): 

98 Which linearity algorithm to use.<br> 

99 - `hc()`: `["hc", "harvey", "harvey-collier"]`<br> 

100 - `lm()`: `["lm", "lagrange", "lagrange-multiplier"]`<br> 

101 - `rb()`: `["rb", "rainbow"]`<br> 

102 - `rr()`: `["rr", "reset", "ramsey-reset"]`<br> 

103 Default: `"rr"` 

104 kwargs (Union[float, int, str, bool, ArrayLike, None]): 

105 Additional keyword arguments passed to the underlying test function. 

106 

107 Raises: 

108 (ValueError): 

109 When the given value for `algorithm` is not valid. 

110 

111 Returns: 

112 (Union[tuple[float, float], tuple[float, float, float, float], ContrastResults]): 

113 - For `"hc"`: `(statistic, pvalue)` 

114 - For `"lm"`: `(lm_stat, lm_pval, f_stat, f_pval)` 

115 - For `"rb"`: `(statistic, pvalue)` 

116 - For `"rr"`: `ContrastResults` object 

117 

118 !!! success "Credit" 

119 Calculations are performed by `statsmodels`. 

120 

121 ???+ example "Examples" 

122 

123 ```pycon {.py .python linenums="1" title="Setup"} 

124 >>> import statsmodels.api as sm 

125 >>> from ts_stat_tests.linearity.tests import linearity 

126 >>> from ts_stat_tests.utils.data import data_line, data_random 

127 >>> x = sm.add_constant(data_line) 

128 >>> y = 3 + 2 * data_line + data_random 

129 >>> res = sm.OLS(y, x).fit() 

130 

131 ``` 

132 

133 ```pycon {.py .python linenums="1" title="Example 1: Ramsey RESET test"} 

134 >>> result = linearity(res, algorithm="rr") 

135 >>> print(f"p-value: {result.pvalue:.4f}") 

136 p-value: 0.9912 

137 

138 ``` 

139 

140 ```pycon {.py .python linenums="1" title="Example 2: Rainbow test"} 

141 >>> stat, pval = linearity(res, algorithm="rb") 

142 >>> print(f"Rainbow p-value: {pval:.4f}") 

143 Rainbow p-value: 0.5391 

144 

145 ``` 

146 """ 

147 options: dict[str, tuple[str, ...]] = { 

148 "hc": ("hc", "harvey", "harvey-collier"), 

149 "lm": ("lm", "lagrange", "lagrange-multiplier"), 

150 "rb": ("rb", "rainbow"), 

151 "rr": ("rr", "reset", "ramsey-reset"), 

152 } 

153 

154 # Internal helper to handle kwargs casting for ty 

155 def _call( 

156 func: Callable[..., Union[tuple[float, ...], ContrastResults]], 

157 **args: Union[ 

158 float, 

159 int, 

160 str, 

161 bool, 

162 ArrayLike, 

163 None, 

164 RegressionResults, 

165 RegressionResultsWrapper, 

166 ], 

167 ) -> Union[tuple[float, ...], ContrastResults]: 

168 """ 

169 !!! note "Summary" 

170 Internal helper to handle keyword arguments types. 

171 

172 Params: 

173 func (Callable[..., Union[tuple[float, ...], ContrastResults]]): 

174 The function to call. 

175 args (Union[float, int, str, bool, ArrayLike, None, RegressionResults, RegressionResultsWrapper]): 

176 The arguments to pass. 

177 

178 Returns: 

179 (Union[tuple[float, ...], ContrastResults]): 

180 The result of the function call. 

181 

182 ???+ example "Examples" 

183 ```python 

184 # Internal use only 

185 ``` 

186 """ 

187 return func(**args) 

188 

189 if algorithm in options["hc"]: 

190 return _call(_hc, res=res, **kwargs) 

191 

192 if algorithm in options["lm"]: 

193 return _call(_lm, resid=res.resid, exog=res.model.exog, **kwargs) 

194 

195 if algorithm in options["rb"]: 

196 return _call(_rb, res=res, **kwargs) 

197 

198 if algorithm in options["rr"]: 

199 return _call(_rr, res=res, **kwargs) 

200 

201 raise ValueError( 

202 generate_error_message( 

203 parameter_name="algorithm", 

204 value_parsed=algorithm, 

205 options=options, 

206 ) 

207 ) 

208 

209 

210@typechecked 

211def is_linear( 

212 res: Union[RegressionResults, RegressionResultsWrapper], 

213 algorithm: str = "rr", 

214 alpha: float = 0.05, 

215 **kwargs: Union[float, int, str, bool, ArrayLike, None], 

216) -> dict[str, Union[str, float, bool, None]]: 

217 """ 

218 !!! note "Summary" 

219 Test whether the relationship in a fitted model is `linear` or not. 

220 

221 ???+ abstract "Details" 

222 This function returns a dictionary containing the test results and a boolean indicating whether the null hypothesis of linearity is rejected at the given significance level. 

223 

224 Params: 

225 res (Union[RegressionResults, RegressionResultsWrapper]): 

226 The fitted regression model to be checked. 

227 algorithm (str): 

228 Which linearity algorithm to use. See [`linearity()`][ts_stat_tests.linearity.tests.linearity] for options. 

229 Default: `"rr"` 

230 alpha (float): 

231 The significance level for the test. 

232 Default: `0.05` 

233 kwargs (Union[float, int, str, bool, ArrayLike, None]): 

234 Additional keyword arguments passed to the underlying test function. 

235 

236 Returns: 

237 (dict[str, Union[str, float, bool, None]]): 

238 A dictionary with the following keys: 

239 - `"algorithm"` (str): The name of the algorithm used. 

240 - `"statistic"` (float): The test statistic. 

241 - `"pvalue"` (float): The p-value of the test. 

242 - `"alpha"` (float): The significance level used. 

243 - `"result"` (bool): Whether the relationship is linear (i.e., p-value > alpha). 

244 

245 ???+ example "Examples" 

246 

247 ```pycon {.py .python linenums="1" title="Setup"} 

248 >>> import statsmodels.api as sm 

249 >>> from ts_stat_tests.linearity.tests import is_linear 

250 >>> from ts_stat_tests.utils.data import data_line, data_random 

251 >>> x = sm.add_constant(data_line) 

252 >>> y = 3 + 2 * data_line + data_random 

253 >>> res = sm.OLS(y, x).fit() 

254 

255 ``` 

256 

257 ```pycon {.py .python linenums="1" title="Example 1: Check linearity with RR"} 

258 >>> result = is_linear(res, algorithm="rr") 

259 >>> print(result["result"]) 

260 True 

261 

262 ``` 

263 """ 

264 options: dict[str, tuple[str, ...]] = { 

265 "hc": ("hc", "harvey", "harvey-collier"), 

266 "lm": ("lm", "lagrange", "lagrange-multiplier"), 

267 "rb": ("rb", "rainbow"), 

268 "rr": ("rr", "reset", "ramsey-reset"), 

269 } 

270 

271 raw_res = linearity(res=res, algorithm=algorithm, **kwargs) 

272 

273 if algorithm in options["hc"] or algorithm in options["rb"]: 

274 # raw_res is tuple[float, float] 

275 stat, pvalue = raw_res[0], raw_res[1] # type: ignore 

276 elif algorithm in options["lm"]: 

277 # raw_res is tuple[float, float, float, float] 

278 stat, pvalue = raw_res[0], raw_res[1] # type: ignore 

279 else: 

280 # raw_res is ContrastResults (Ramsey RESET) 

281 stat = float(getattr(raw_res, "statistic", getattr(raw_res, "fvalue", np.nan))) 

282 pvalue = float(getattr(raw_res, "pvalue", np.nan)) 

283 

284 return { 

285 "algorithm": algorithm, 

286 "statistic": float(stat), 

287 "pvalue": float(pvalue), 

288 "alpha": alpha, 

289 "result": bool(pvalue > alpha), 

290 }