Sped up covariance estimate via numba

.coveragerc

@@ -29,6 +29,8 @@ exclude_lines =
     cdef
     # Cython functions with void
     cdef void
+    # Numba jit decorators
+    @jit
 
 
 include = */arch/*

arch/compat/numba.py

@@ -1,12 +1,11 @@
 import functools
 import os
-from typing import Any, Callable
+from typing import Any, Callable, Optional
 
 from arch.utility.exceptions import PerformanceWarning
 
 DISABLE_NUMBA = os.environ.get("ARCH_DISABLE_NUMBA", False) in ("1", "true", "True")
 
-
 performance_warning: str = """
 numba is not available, and this function is being executed without JIT
 compilation. Either install numba or reinstalling after installing Cython
@@ -22,14 +21,25 @@
 
 except ImportError:
 
-    def jit(func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
-        def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
-            import warnings
+    def jit(
+        function_or_signature: Optional[Callable[..., Any]] = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Any:
+        if function_or_signature is not None and callable(function_or_signature):
+            return function_or_signature
+
+        def wrap(func):
+            @functools.wraps(func)
+            def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
+                import warnings
+
+                warnings.warn(performance_warning, PerformanceWarning)
+                return func(*args, **kwargs)
 
-            warnings.warn(performance_warning, PerformanceWarning)
-            return func(*args, **kwargs)
+            return wrapper
 
-        return wrapper
+        return wrap
 
 
 __all__ = ["jit", "PerformanceWarning"]

arch/covariance/kernel.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from arch.compat.numba import jit
+
 from abc import ABC, abstractmethod
 from functools import cached_property
 from typing import SupportsInt, cast
@@ -136,6 +138,14 @@ def one_sided_strict(self) -> Float64Array | DataFrame:
         return self._wrap(self._oss)
 
 
+@jit(nopython=True)
+def _cov_jit(df, k, num_weights, w, x):
+    oss = np.zeros((k, k))
+    for i in range(1, num_weights):
+        oss += w[i] * (x[i:].T @ x[:-i]) / df
+    return oss
+
+
 class CovarianceEstimator(ABC):
     r"""
     %(kernel_name)s kernel covariance estimation.
@@ -392,15 +402,14 @@ def cov(self) -> CovarianceEstimate:
         --------
         CovarianceEstimate
         """
-        x = np.asarray(self._x)
+        x = np.ascontiguousarray(self._x)
         k = x.shape[1]
         df = self._df
         sr = x.T @ x / df
         w = self.kernel_weights
         num_weights = w.shape[0]
-        oss = np.zeros((k, k))
-        for i in range(1, num_weights):
-            oss += w[i] * (x[i:].T @ x[:-i]) / df
+
+        oss = _cov_jit(df, k, num_weights, w, x)
 
         labels = self._x_orig.columns if isinstance(self._x_orig, DataFrame) else None
         return CovarianceEstimate(sr, oss, labels)