Improve typing 03 25

arch/bootstrap/_samplers.pyi

@@ -1,5 +1,5 @@
-from arch.typing import Float64Array, Int64Array
+from arch.typing import Float64Array, Int64Array1D
 
 def stationary_bootstrap_sample(
-    indices: Int64Array, u: Float64Array, p: float
-) -> Int64Array: ...
+    indices: Int64Array1D, u: Float64Array, p: float
+) -> Int64Array1D: ...

arch/bootstrap/base.py

@@ -15,6 +15,7 @@
     BootstrapIndexT,
     Float64Array,
     Int64Array,
+    Int64Array1D,
     Literal,
     NDArray,
     RandomStateState,
@@ -55,12 +56,12 @@ def _get_prng_state(
 
 def _get_random_integers(
     prng: Union[Generator, RandomState], upper: int, *, size: int = 1
-) -> Int64Array:
+) -> Int64Array1D:
     if isinstance(prng, Generator):
-        return prng.integers(upper, size=size, dtype=np.int64)
+        return cast(Int64Array1D, prng.integers(upper, size=size, dtype=np.int64))
     else:
         assert isinstance(prng, RandomState)
-        return prng.randint(upper, size=size, dtype=np.int64)
+        return cast(Int64Array1D, prng.randint(upper, size=size, dtype=np.int64))
 
 
 def _single_optimal_block(x: Float64Array) -> tuple[float, float]:
@@ -438,7 +439,7 @@ def __init__(
                     raise ValueError(
                         "All inputs must have the same number of elements in axis 0"
                     )
-        self._index: BootstrapIndexT = np.arange(self._num_items)
+        self._index: BootstrapIndexT = np.arange(self._num_items, dtype=np.int64)
 
         self._parameters: list[int] = []
         self.pos_data: tuple[Union[AnyArray, pd.Series, pd.DataFrame], ...] = args
@@ -626,7 +627,7 @@ def reset(self, use_seed: bool = True) -> None:
             False or if no seed has been set, the bootstrap will be reset
             to the initial state.  Default is True
         """
-        self._index = np.arange(self._num_items)
+        self._index = np.arange(self._num_items, dtype=np.int64)
         self._resample()
         self.state = self._initial_state
 
@@ -1380,7 +1381,7 @@ def __init__(
 
     def update_indices(
         self,
-    ) -> tuple[list[Int64Array], dict[str, Int64Array]]:
+    ) -> tuple[list[Int64Array1D], dict[str, Int64Array1D]]:
         """
         Update indices for the next iteration of the bootstrap.  This must
         be overridden when creating new bootstraps.
@@ -1424,8 +1425,14 @@ def reset(self, use_seed: bool = True) -> None:
             False or if no seed has been set, the bootstrap will be reset
             to the initial state.  Default is True
         """
-        pos_indices = [np.arange(self._num_arg_items[i]) for i in range(self._num_args)]
-        kw_indices = {key: np.arange(self._num_kw_items[key]) for key in self._kwargs}
+        pos_indices: list[Int64Array1D] = [
+            np.arange(self._num_arg_items[i], dtype=np.int64)
+            for i in range(self._num_args)
+        ]
+        kw_indices: dict[str, Int64Array1D] = {
+            key: np.arange(self._num_kw_items[key], dtype=np.int64)
+            for key in self._kwargs
+        }
         self._index = pos_indices, kw_indices
         self._resample()
         self.state = self._initial_state
@@ -1595,21 +1602,21 @@ def _repr_html(self) -> str:
         html += ", <strong>ID</strong>: " + hex(id(self)) + ")"
         return html
 
-    def update_indices(self) -> Int64Array:
+    def update_indices(self) -> Int64Array1D:
         num_blocks = self._num_items // self.block_size
         if num_blocks * self.block_size < self._num_items:
             num_blocks += 1
         indices = _get_random_integers(
             self._generator, self._num_items, size=num_blocks
         )
-        indices = indices[:, None] + np.arange(self.block_size)
-        indices = indices.flatten()
+        _indices = indices[:, None] + np.arange(self.block_size, dtype=np.int64)
+        indices = _indices.flatten()
         indices %= self._num_items
 
         if indices.shape[0] > self._num_items:
-            return indices[: self._num_items]
+            return cast(Int64Array1D, indices[: self._num_items])
         else:
-            return indices
+            return cast(Int64Array1D, indices)
 
 
 class StationaryBootstrap(CircularBlockBootstrap):
@@ -1707,7 +1714,7 @@ def __init__(
         )
         self._p = 1.0 / block_size
 
-    def update_indices(self) -> Int64Array:
+    def update_indices(self) -> Int64Array1D:
         indices = _get_random_integers(
             self._generator, self._num_items, size=self._num_items
         )
@@ -1816,7 +1823,7 @@ def __init__(
             block_size, *args, random_state=random_state, seed=seed, **kwargs
         )
 
-    def update_indices(self) -> Int64Array:
+    def update_indices(self) -> Int64Array1D:
         num_blocks = self._num_items // self.block_size
         if num_blocks * self.block_size < self._num_items:
             num_blocks += 1
@@ -1826,7 +1833,7 @@ def update_indices(self) -> Int64Array:
         indices = indices.flatten()
 
         if indices.shape[0] > self._num_items:
-            return indices[: self._num_items]
+            return cast(Int64Array1D, indices[: self._num_items])
         else:
             return indices
 
@@ -1846,6 +1853,6 @@ def __init__(
     def update_indices(
         self,
     ) -> Union[
-        Int64Array, tuple[list[Int64Array], dict[str, Int64Array]]
+        Int64Array1D, tuple[list[Int64Array1D], dict[str, Int64Array1D]]
     ]:  # pragma: no cover
         raise NotImplementedError

arch/bootstrap/multiple_comparison.py

@@ -16,7 +16,7 @@
     ArrayLike2D,
     BoolArray,
     Float64Array,
-    Int64Array,
+    Int64Array1D,
     IntArray,
     Literal,
     Uint32Array,
@@ -448,7 +448,7 @@ def __init__(
         self.k: int = self.models.shape[1]
         self.reps: int = reps
         self.size: float = size
-        self._superior_models: Optional[list[Hashable]] = None
+        self._superior_models: Optional[list[int]] = None
         self.bootstrap: CircularBlockBootstrap = self.spa.bootstrap
 
         self._model = "StepM"
@@ -472,16 +472,13 @@ def compute(self) -> None:
         # 1. Run SPA
         self.spa.compute()
         # 2. If any models superior, store indices, remove and re-run SPA
-        better_models = list(self.spa.better_models(self.size))
+        better_models = list(int(i) for i in self.spa.better_models(self.size))
         all_better_models = better_models[:]
         # 3. Stop if nothing superior
         while better_models and (len(better_models) < self.k):
             # A. Use Selector to remove better models
             selector = np.ones(self.k, dtype=np.bool_)
-            if isinstance(self.models, pd.DataFrame):  # Columns
-                selector[self.models.columns.isin(all_better_models)] = False
-            else:
-                selector[np.array(all_better_models)] = False
+            selector[np.array(all_better_models)] = False
             self.spa.subset(selector)
             # B. Rerun
             self.spa.compute()
@@ -494,7 +491,7 @@ def compute(self) -> None:
         self._superior_models = all_better_models
 
     @property
-    def superior_models(self) -> list[Hashable]:
+    def superior_models(self) -> Union[list[int], Sequence[Hashable]]:
         """
         List of the indices or column names of the superior models
 
@@ -507,6 +504,8 @@ def superior_models(self) -> list[Hashable]:
         if self._superior_models is None:
             msg = "compute must be called before accessing superior_models"
             raise RuntimeError(msg)
+        if isinstance(self.models, pd.DataFrame):
+            return list(self.models.columns[self._superior_models])
         return self._superior_models
 
 
@@ -781,7 +780,7 @@ def better_models(
         self,
         pvalue: float = 0.05,
         pvalue_type: Literal["lower", "consistent", "upper"] = "consistent",
-    ) -> Union[Int64Array, list[Hashable]]:
+    ) -> Union[Int64Array1D]:
         """
         Returns set of models rejected as being equal-or-worse than the
         benchmark
@@ -811,10 +810,7 @@ def better_models(
         crit_val = self.critical_values(pvalue=pvalue)[pvalue_type]
         better_models = self._loss_diff.mean(0) > crit_val
         better_models = np.logical_and(better_models, self._selector)
-        if isinstance(self.models, pd.DataFrame):
-            return list(self.models.columns[better_models])
-        else:
-            return np.argwhere(better_models).flatten()
+        return np.argwhere(better_models).flatten()
 
     def _check_compute(self) -> None:
         if self._pvalues:

arch/compat/numba.py

@@ -4,6 +4,13 @@
 
 from arch.utility.exceptions import PerformanceWarning
 
+try:
+    import numba  # noqa
+
+    HAS_NUMBA = True
+except ImportError:
+    HAS_NUMBA = False
+
 DISABLE_NUMBA = os.environ.get("ARCH_DISABLE_NUMBA", False) in ("1", "true", "True")
 
 performance_warning: str = """
@@ -51,4 +58,4 @@ def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
         return wrap
 
 
-__all__ = ["jit", "PerformanceWarning"]
+__all__ = ["jit", "PerformanceWarning", "HAS_NUMBA", "DISABLE_NUMBA"]

arch/covariance/kernel.py

@@ -1,4 +1,4 @@
-from arch.compat.numba import jit
+from arch.compat.numba import DISABLE_NUMBA, HAS_NUMBA, jit
 
 from abc import ABC, abstractmethod
 from functools import cached_property
@@ -136,8 +136,7 @@ def one_sided_strict(self) -> Union[Float64Array, DataFrame]:
         return self._wrap(self._oss)
 
 
-@jit(nopython=True)
-def _cov_jit(
+def _cov_kernel(
     df: int, k: int, num_weights: int, w: np.ndarray, x: np.ndarray
 ) -> np.ndarray:
     oss = np.zeros((k, k))
@@ -146,6 +145,12 @@ def _cov_jit(
     return oss
 
 
+if HAS_NUMBA and not DISABLE_NUMBA:
+    _cov_jit = jit(_cov_kernel, nopython=True)
+else:
+    _cov_jit = _cov_kernel
+
+
 class CovarianceEstimator(ABC):
     r"""
     %(kernel_name)s kernel covariance estimation.

arch/tests/unitroot/cointegration_data.py

@@ -1,10 +1,12 @@
 from arch.compat.pandas import MONTH_END
 
+from typing import cast
+
 import numpy as np
 import pandas as pd
 import pytest
 
-from arch.typing import ArrayLike2D, Float64Array
+from arch.typing import ArrayLike2D, Float64Array, Float64Array2D
 
 
 @pytest.fixture(scope="module", params=[True, False])
@@ -65,6 +67,8 @@ def trivariate_data(request) -> tuple[ArrayLike2D, ArrayLike2D]:
         dt_index = pd.date_range("1-1-2000", periods=nobs, freq=MONTH_END)
         cols = [f"y{i}" for i in range(1, 4)]
         data = pd.DataFrame(y, columns=cols, index=dt_index)
-        return data.iloc[:, :1], data.iloc[:, 1:]
+        return cast(pd.DataFrame, data.iloc[:, :1]), cast(
+            pd.DataFrame, data.iloc[:, 1:]
+        )
 
-    return y[:, :1], y[:, 1:]
+    return cast(Float64Array2D, y[:, :1]), cast(Float64Array2D, y[:, 1:])

arch/tests/univariate/test_arch_in_mean.py

@@ -17,10 +17,11 @@
 
 SP500 = 100 * sp500.load()["Adj Close"].pct_change().dropna()
 SP500 = SP500.iloc[SP500.shape[0] // 2 :]
-X = pd.concat([SP500, SP500], axis=1)
-X.columns = pd.Index([0, 1])
 RANDOMSTATE = np.random.RandomState(12349876)
-X.loc[:, :] = RANDOMSTATE.standard_normal(X.shape)
+X = pd.DataFrame(
+    RANDOMSTATE.standard_normal((SP500.shape[0], 2)), columns=[0, 1], index=SP500.index
+)
+
 
 SUPPORTED = [
     HARCH,

arch/tests/utility/test_array.py

@@ -18,6 +18,7 @@
     ensure2d,
     find_index,
     parse_dataframe,
+    to_array_1d,
 )
 
 
@@ -295,3 +296,29 @@ class Dummy(metaclass=ConcreteClassMeta):
             @abstractmethod
             def func(self):
                 pass
+
+
+@pytest.mark.parametrize(
+    "arr",
+    [
+        np.array([1.0], dtype=float),
+        np.array(0),
+        np.array([0, 1]),
+        np.array([[2, 3, 4]]),
+        Series([1, 2, 3]),
+    ],
+)
+def test_to_array_1d(arr):
+    converted = to_array_1d(arr)
+    assert isinstance(converted, np.ndarray)
+    assert converted.ndim == 1
+    assert converted.dtype == np.float64
+
+
+def test_to_array_1d_err():
+    with pytest.raises(ValueError, match="x must be 1D"):
+        to_array_1d(np.array([[1, 2], [3, 4]]))
+    with pytest.raises(TypeError, match="x must be a Series or ndarray"):
+        to_array_1d(0)
+    with pytest.raises(TypeError, match="x must be a Series or ndarray"):
+        to_array_1d(DataFrame([[0, 1]]))

arch/typing.py

@@ -38,17 +38,21 @@
 Float64Array = np.ndarray[tuple[int, ...], np.dtype[np.float64]]  # pragma: no cover
 Float64Array1D = np.ndarray[tuple[int], np.dtype[np.float64]]  # pragma: no cover
 Float64Array2D = np.ndarray[tuple[int, int], np.dtype[np.float64]]  # pragma: no cover
-Int64Array = np.ndarray[tuple[int, ...], np.dtype[np.longlong]]  # pragma: no cover
-Int64Array2D = np.ndarray[tuple[int, int], np.dtype[np.longlong]]  # pragma: no cover
+Int64Array = np.ndarray[tuple[int, ...], np.dtype[np.int64]]  # pragma: no cover
+Int64Array1D = np.ndarray[tuple[int], np.dtype[np.int64]]  # pragma: no cover
+#     tuple[list[np.ndarray[tuple[int], np.dtype[np.signedinteger[Any]]]], dict[str, np.ndarray[tuple[int], np.dtype[np.signedinteger[Any]]]]]
+Int64Array2D = np.ndarray[tuple[int, int], np.dtype[np.int64]]  # pragma: no cover
 Int32Array = np.ndarray[tuple[int, ...], np.dtype[np.intc]]  # pragma: no cover
-IntArray = np.ndarray[tuple[int, ...], np.dtype[np.int_]]  # pragma: no cover
+IntArray = np.ndarray[tuple[int, ...], np.dtype[np.int64]]  # pragma: no cover
 BoolArray = np.ndarray[tuple[int, ...], np.dtype[np.bool_]]  # pragma: no cover
 AnyArray = np.ndarray[tuple[int, ...], Any]  # pragma: no cover
 AnyArray1D = np.ndarray[tuple[int], Any]  # pragma: no cover
 Uint32Array = np.ndarray[tuple[int, ...], np.dtype[np.uintc]]  # pragma: no cover
 
 BootstrapIndexT = Union[
-    Int64Array, tuple[Int64Array, ...], tuple[list[Int64Array], dict[str, Int64Array]]
+    Int64Array1D,
+    tuple[Int64Array1D, ...],
+    tuple[list[Int64Array1D], dict[str, Int64Array1D]],
 ]
 RandomStateState = tuple[str, Uint32Array, int, int, float]
 

arch/unitroot/critical_values/simulation/kpss_critical_values_simulation.py

@@ -10,6 +10,7 @@
 from numpy.random import RandomState
 import pandas as pd
 
+from arch.typing import Float64Array
 from arch.utility.timeseries import add_trend
 
 
@@ -30,7 +31,7 @@ def simulate_kpss(
     standard_normal = rng.standard_normal
 
     e = standard_normal((nobs, b))
-    z = np.ones((nobs, 1))
+    z: Float64Array = np.ones((nobs, 1))
     if trend == "ct":
         z = add_trend(z, trend="t")
     zinv = np.linalg.pinv(z)