From f1bf2a6534ac28cf800dcafd70fe27f8337fc01a Mon Sep 17 00:00:00 2001 From: mmcky Date: Thu, 7 Mar 2024 11:24:29 +1100 Subject: [PATCH] TST: update arrays variable to contain numpy arrays --- lectures/ifp_egm.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/lectures/ifp_egm.md b/lectures/ifp_egm.md index c3116899..4f17e992 100644 --- a/lectures/ifp_egm.md +++ b/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: