feat: Add JIT acceleration for Kalman filter and enhance system building functions

Author: pablo-reyes8

README.md

@@ -6,6 +6,8 @@ This repository captures the core building blocks of a Dynare-style DSGE workflo
 
 ```bash
 python -m pip install -e ".[dev]"
+# Optional JIT acceleration for the Kalman loop:
+python -m pip install -e ".[speed]"
 python -m pytest
 python scripts/run_pipeline.py --config configs/tiny_ar1.yaml
 python scripts/run_pipeline.py --config configs/nk_full_yaml.yaml --dry-run

pyproject.toml

@@ -23,6 +23,9 @@ dev = [
     "pytest>=7.4",
     "pytest-cov>=4.1",
 ]
+speed = [
+    "numba>=0.59",
+]
 
 [tool.setuptools.packages.find]
 where = ["."]

src/analysis/impulse_responses.py

@@ -1,5 +1,5 @@
-import numpy as np 
-from src.inference.likelihoods import *
+import numpy as np 
+from src.inference.likelihoods import *
 
 def compute_irfs(
     draws_work,
@@ -16,22 +16,22 @@ def compute_irfs(
     burn_in=0,
     quantiles=(0.16, 0.5, 0.84),
     shock_indices=None,
-    observable_names=None,
-    shock_names=None,
-    shock_scale="std",
-    div=0.0,
-    steady=None,):
+    observable_names=None,
+    shock_names=None,
+    shock_scale="std",
+    div=0.0,
+    steady=None,):
 
     draws = np.asarray(draws_work, dtype=float)
     draws = draws[burn_in:] if burn_in > 0 else draws
 
-    if horizon < 1:
-        raise ValueError("horizon debe ser al menos 1.")
-
-    n_shocks = len(eps_t)
-    shock_indices = range(n_shocks) if shock_indices is None else shock_indices
-    if shock_scale not in {"std", "unit"}:
-        raise ValueError("shock_scale debe ser 'std' o 'unit'.")
+    if horizon < 1:
+        raise ValueError("horizon debe ser al menos 1.")
+
+    n_shocks = len(eps_t)
+    shock_indices = range(n_shocks) if shock_indices is None else shock_indices
+    if shock_scale not in {"std", "unit"}:
+        raise ValueError("shock_scale debe ser 'std' o 'unit'.")
 
     if shock_names is None:
         shock_names = [str(sym) for sym in eps_t]
@@ -41,24 +41,27 @@ def compute_irfs(
     Psi2_example = None
     valid_draws = []
     for theta in draws:
-        Theta1, C, Theta0, eu, Psi0, Psi2 = st_sp(theta,
-                                    equations,y_t,
-                                    y_tp1,
-                                    eps_t,
+        Theta1, C, Theta0, eu, Psi0, Psi2 = st_sp(theta,
+                                    equations,y_t,
+                                    y_tp1,
+                                    eps_t,
                                     registry,
-                                    y_tm1=y_tm1,
-                                    eta_t=eta_t,
-                                    measurement=measurement,
-                                    steady=steady,
-                                    div=div)
+                                    y_tm1=y_tm1,
+                                    eta_t=eta_t,
+                                    measurement=measurement,
+                                    steady=steady,
+                                    div=div)
 
         if eu[0] < 1 or eu[1] < 1:
             continue
-        try:
-            Q = registry.build_Q(theta, n_shocks)
-            L = np.linalg.cholesky(Q)
-        except np.linalg.LinAlgError:
-            continue
+        if shock_scale == "std":
+            try:
+                Q = registry.build_Q(theta, n_shocks)
+                L = np.linalg.cholesky(Q)
+            except np.linalg.LinAlgError:
+                continue
+        else:
+            L = None
         valid_draws.append((theta, Theta1, Theta0, Psi0, Psi2, L))
         if Psi2_example is None:
             Psi2_example = Psi2
@@ -74,23 +77,23 @@ def compute_irfs(
 
     irf_store = {shock_names[j]: [] for j in shock_indices}
 
-    for theta, Theta1, Theta0, Psi0, Psi2, L in valid_draws:
-        for j in shock_indices:
-            eps_path = np.zeros((n_shocks, horizon))
-            if shock_scale == "std":
-                # Shock de una desviacion estandar, usando la covarianza estimada Q.
-                eps_path[:, 0] = L[:, j]
-            else:
-                # Shock estructural unitario, como stoch_simul tras fijar var eps_j = 1.
-                eps_path[j, 0] = 1.0
-
-            state = np.zeros((Theta1.shape[0], horizon))
-            obs = np.zeros((n_obs, horizon))
-
-            for h in range(horizon):
-                previous = state[:, h - 1] if h > 0 else np.zeros(Theta1.shape[0])
-                state[:, h] = Theta1 @ previous + Theta0 @ eps_path[:, h]
-                obs[:, h] = Psi2 @ state[:, h]
+    for theta, Theta1, Theta0, Psi0, Psi2, L in valid_draws:
+        for j in shock_indices:
+            eps_path = np.zeros((n_shocks, horizon))
+            if shock_scale == "std":
+                # Shock de una desviacion estandar, usando la covarianza estimada Q.
+                eps_path[:, 0] = L[:, j]
+            else:
+                # Shock estructural unitario, como stoch_simul tras fijar var eps_j = 1.
+                eps_path[j, 0] = 1.0
+
+            state = np.zeros((Theta1.shape[0], horizon))
+            obs = np.zeros((n_obs, horizon))
+
+            for h in range(horizon):
+                previous = state[:, h - 1] if h > 0 else np.zeros(Theta1.shape[0])
+                state[:, h] = Theta1 @ previous + Theta0 @ eps_path[:, h]
+                obs[:, h] = Psi2 @ state[:, h]
 
             irf_store[shock_names[j]].append(obs)
 
@@ -100,19 +103,19 @@ def compute_irfs(
 
         arr = np.stack(paths, axis=0) 
         q_values = np.quantile(arr, quant_array, axis=0)  
-        results[shock_name] = {
-            "observables": observable_names,
-            "horizon": horizon,
-            "quantiles": quant_array,
-            "shock_scale": shock_scale,
-            "summary": q_values,
-            "raw": arr,}
+        results[shock_name] = {
+            "observables": observable_names,
+            "horizon": horizon,
+            "quantiles": quant_array,
+            "shock_scale": shock_scale,
+            "summary": q_values,
+            "raw": arr,}
 
     return results
 
 
-def plot_irf_bands(irf_dict, shocks=None, start=0, figsize=(7, 4), colors=None):
-    import matplotlib.pyplot as plt
+def plot_irf_bands(irf_dict, shocks=None, start=0, figsize=(7, 4), colors=None):
+    import matplotlib.pyplot as plt
 
     shocks = shocks or list(irf_dict.keys())
     colors = colors or {
@@ -152,4 +155,4 @@ def plot_irf_bands(irf_dict, shocks=None, start=0, figsize=(7, 4), colors=None):
 
         axes[-1].set_xlabel("Horizonte")
         plt.tight_layout()
-        plt.show()
+        plt.show()

src/inference/kalman_fast.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+try:
+    from numba import njit
+
+    HAS_NUMBA = True
+except Exception:  # pragma: no cover - depends on optional dependency
+    njit = None
+    HAS_NUMBA = False
+
+
+if HAS_NUMBA:  # pragma: no cover - exercised only when numba is installed
+    import numpy as np
+
+    @njit(cache=True)
+    def _kalman_loglike_numba(y, Theta1, C, Theta0, Q, H, Psi0, Psi2, s_bar, P_bar):
+        T = y.shape[0]
+        n_y = y.shape[1]
+        n_s = Theta1.shape[0]
+        const = n_y * np.log(2.0 * np.pi)
+        ll_sum = 0.0
+        eye = np.eye(n_s)
+
+        for t in range(T):
+            s_hat = C + Theta1 @ s_bar
+            P_hat = Theta1 @ P_bar @ Theta1.T + Theta0 @ Q @ Theta0.T
+            P_hat = 0.5 * (P_hat + P_hat.T)
+
+            e_t = y[t, :] - Psi0 - Psi2 @ s_hat
+            S_t = Psi2 @ P_hat @ Psi2.T + H
+
+            try:
+                L = np.linalg.cholesky(S_t)
+            except Exception:
+                return -np.inf
+
+            logdet = 0.0
+            for j in range(n_y):
+                logdet += 2.0 * np.log(L[j, j])
+
+            z = np.linalg.solve(L, e_t)
+            Sinv_e = np.linalg.solve(L.T, z)
+
+            ZP = Psi2 @ P_hat.T
+            tmp = np.linalg.solve(L, ZP)
+            K = np.linalg.solve(L.T, tmp).T
+
+            s_bar = s_hat + K @ e_t
+            I_KZ = eye - K @ Psi2
+            P_bar = I_KZ @ P_hat @ I_KZ.T + K @ H @ K.T
+            P_bar = 0.5 * (P_bar + P_bar.T)
+
+            ll_sum += -0.5 * (const + logdet + e_t @ Sinv_e)
+
+        return ll_sum
+
+
+def kalman_loglike_numba(y, Theta1, C, Theta0, Q, H, Psi0, Psi2, s_bar, P_bar):
+    if not HAS_NUMBA:
+        return None
+    return float(_kalman_loglike_numba(y, Theta1, C, Theta0, Q, H, Psi0, Psi2, s_bar, P_bar))