TST: update arrays variable to contain numpy arrays

lectures/ifp_egm.md

@@ -416,7 +416,7 @@ model and run the cross-check.
 
 ```{code-cell}
 @numba.jit
-def K_egm_nb(a_in, σ_in, constants, sizes, s_grid, y_grid, P):
+def K_egm_nb(a_in, σ_in, constants, sizes, arrays):
     """
     The operator K using Numba.
 
@@ -425,7 +425,7 @@ def K_egm_nb(a_in, σ_in, constants, sizes, s_grid, y_grid, P):
     # Simplify names
     β, R, γ = constants
     s_size, y_size = sizes
-    # s_grid, y_grid, P = arrays
+    s_grid, y_grid, P = arrays
 
     def u_prime(c):
         return c**(-γ)
@@ -466,7 +466,8 @@ def successive_approx_numba(model,        # Class with model information
     constants, sizes, arrays = model
     s_size, y_size = sizes
     # make NumPy versions of arrays
-    s_grid, y_grid, P = tuple(map(np.array, arrays))  
+    arrays = tuple(map(np.array, arrays))
+    s_grid, y_grid, P = arrays
     
     σ_init = np.repeat(s_grid, y_size)
     σ_init = np.reshape(σ_init, (s_size, y_size))
@@ -478,7 +479,7 @@ def successive_approx_numba(model,        # Class with model information
     error = tol + 1
 
     while i < max_iter and error > tol:
-        a_new, σ_new = K_egm_nb(a_vec, σ_vec, constants, sizes, s_grid, y_grid, P)
+        a_new, σ_new = K_egm_nb(a_vec, σ_vec, constants, sizes, arrays)
         error = np.max(np.abs(σ_vec - σ_new))
         i += 1
         if verbose and i % print_skip == 0: