MAINT: update @jit(nopython=True) to @jit with numba>=0.59 (#395)

* MAINT: update @jit(nopython=True) to @jit with numba>=0.59 * update @njit to @jit * update imports njit to jit and usage * minor fixes for import * remove njit from cass_koopmans_1 * remove njit from ak2
QuantEcon · Nov 18, 2024 · 9ec9e13 · 9ec9e13
1 parent bc7a5cb
commit 9ec9e13
Show file tree

Hide file tree

Showing 31 changed files with 150 additions and 150 deletions.
diff --git a/lectures/aiyagari.md b/lectures/aiyagari.md
@@ -283,7 +283,7 @@ class Household:
 
 # Do the hard work using JIT-ed functions
 
-@jit(nopython=True)
+@jit
 def populate_R(R, a_size, z_size, a_vals, z_vals, r, w):
     n = a_size * z_size
     for s_i in range(n):
@@ -297,7 +297,7 @@ def populate_R(R, a_size, z_size, a_vals, z_vals, r, w):
             if c > 0:
                 R[s_i, new_a_i] = np.log(c)  # Utility
 
-@jit(nopython=True)
+@jit
 def populate_Q(Q, a_size, z_size, Π):
     n = a_size * z_size
     for s_i in range(n):
@@ -307,7 +307,7 @@ def populate_Q(Q, a_size, z_size, Π):
                 Q[s_i, a_i, a_i*z_size + next_z_i] = Π[z_i, next_z_i]
 
 
-@jit(nopython=True)
+@jit
 def asset_marginal(s_probs, a_size, z_size):
     a_probs = np.zeros(a_size)
     for a_i in range(a_size):

diff --git a/lectures/ak2.md b/lectures/ak2.md
@@ -408,7 +408,7 @@ $$
 ```{code-cell} ipython3
 import numpy as np
 import matplotlib.pyplot as plt
-from numba import njit
+from numba import jit
 from quantecon.optimize import brent_max
 ```
 
@@ -433,22 +433,22 @@ Knowing $\hat K$, we can calculate other equilibrium objects.
 Let's first define  some Python helper functions.
 
 ```{code-cell} ipython3
-@njit
+@jit
 def K_to_Y(K, α):
 
     return K ** α
 
-@njit
+@jit
 def K_to_r(K, α):
 
     return α * K ** (α - 1)
 
-@njit
+@jit
 def K_to_W(K, α):
 
     return (1 - α) * K ** α
 
-@njit
+@jit
 def K_to_C(K, D, τ, r, α, β):
 
     # optimal consumption for the old when δ=0
@@ -913,7 +913,7 @@ Let's implement this "guess and verify" approach
 We start by defining the Cobb-Douglas utility function
 
 ```{code-cell} ipython3
-@njit
+@jit
 def U(Cy, Co, β):
 
     return (Cy ** β) * (Co ** (1-β))
@@ -924,7 +924,7 @@ We use `Cy_val` to compute the lifetime value of an arbitrary consumption plan,
 Note that it requires knowing future prices $r_{t+1}$ and tax rate $\tau_{t+1}$.
 
 ```{code-cell} ipython3
-@njit
+@jit
 def Cy_val(Cy, W, r_next, τ, τ_next, δy, δo_next, β):
 
     # Co given by the budget constraint

diff --git a/lectures/career.md b/lectures/career.md
@@ -51,7 +51,7 @@ import matplotlib.pyplot as plt
 plt.rcParams["figure.figsize"] = (11, 5)  #set default figure size
 import numpy as np
 import quantecon as qe
-from numba import njit, prange
+from numba import jit, prange
 from quantecon.distributions import BetaBinomial
 from scipy.special import binom, beta
 from mpl_toolkits.mplot3d.axes3d import Axes3D
@@ -234,7 +234,7 @@ def operator_factory(cw, parallel_flag=True):
     F_probs, G_probs = cw.F_probs, cw.G_probs
     F_mean, G_mean = cw.F_mean, cw.G_mean
 
-    @njit(parallel=parallel_flag)
+    @jit(parallel=parallel_flag)
     def T(v):
         "The Bellman operator"
 
@@ -249,7 +249,7 @@ def operator_factory(cw, parallel_flag=True):
 
         return v_new
 
-    @njit
+    @jit
     def get_greedy(v):
         "Computes the v-greedy policy"
 
@@ -472,7 +472,7 @@ T, get_greedy = operator_factory(cw)
 v_star = solve_model(cw, verbose=False)
 greedy_star = get_greedy(v_star)
 
-@njit
+@jit
 def passage_time(optimal_policy, F, G):
     t = 0
     i = j = 0
@@ -485,7 +485,7 @@ def passage_time(optimal_policy, F, G):
             i, j  = qe.random.draw(F), qe.random.draw(G)
         t += 1
 
-@njit(parallel=True)
+@jit(parallel=True)
 def median_time(optimal_policy, F, G, M=25000):
     samples = np.empty(M)
     for i in prange(M):

diff --git a/lectures/cass_koopmans_1.md b/lectures/cass_koopmans_1.md
@@ -4,7 +4,7 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.16.1
+    jupytext_version: 1.16.4
 kernelspec:
   display_name: Python 3 (ipykernel)
   language: python
@@ -66,7 +66,7 @@ Let's start with some standard imports:
 
 ```{code-cell} ipython3
 import matplotlib.pyplot as plt
-from numba import njit, float64
+from numba import jit, float64
 from numba.experimental import jitclass
 import numpy as np
 from quantecon.optimize import brentq
@@ -525,7 +525,7 @@ planning problem.
 $c_0$ instead of $\mu_0$ in the following code.)
 
 ```{code-cell} ipython3
-@njit
+@jit
 def shooting(pp, c0, k0, T=10):
     '''
     Given the initial condition of capital k0 and an initial guess
@@ -610,7 +610,7 @@ When $K_{T+1}$ gets close enough to $0$ (i.e., within an error
 tolerance bounds), we stop.
 
 ```{code-cell} ipython3
-@njit
+@jit
 def bisection(pp, c0, k0, T=10, tol=1e-4, max_iter=500, k_ter=0, verbose=True):
 
     # initial boundaries for guess c0
@@ -804,7 +804,7 @@ over time.
 Let's calculate and  plot the saving rate.
 
 ```{code-cell} ipython3
-@njit
+@jit
 def saving_rate(pp, c_path, k_path):
     'Given paths of c and k, computes the path of saving rate.'
     production = pp.f(k_path[:-1])
@@ -912,7 +912,7 @@ $$ (eq:tildeC)
 A positive fixed point  $C = \tilde C(K)$ exists only if $f\left(K\right)+\left(1-\delta\right)K-f^{\prime-1}\left(\frac{1}{\beta}-\left(1-\delta\right)\right)>0$
 
 ```{code-cell} ipython3
-@njit
+@jit
 def C_tilde(K, pp):
 
     return pp.f(K) + (1 - pp.δ) * K - pp.f_prime_inv(1 / pp.β - 1 + pp.δ)
@@ -931,11 +931,11 @@ K = \tilde K(C)
 $$ (eq:tildeK)
 
 ```{code-cell} ipython3
-@njit
+@jit
 def K_diff(K, C, pp):
     return pp.f(K) - pp.δ * K - C
 
-@njit
+@jit
 def K_tilde(C, pp):
 
     res = brentq(K_diff, 1e-6, 100, args=(C, pp))
@@ -951,7 +951,7 @@ It is thus the intersection of the two curves $\tilde{C}$ and $\tilde{K}$ that w
 We can compute $K_s$ by solving the equation $K_s = \tilde{K}\left(\tilde{C}\left(K_s\right)\right)$
 
 ```{code-cell} ipython3
-@njit
+@jit
 def K_tilde_diff(K, pp):
 
     K_out = K_tilde(C_tilde(K, pp), pp)
@@ -1003,7 +1003,7 @@ In addition to the three curves, Figure {numref}`stable_manifold` plots  arrows
 ---
 mystnb:
   figure:
-    caption: "Stable Manifold and Phase Plane"
+    caption: Stable Manifold and Phase Plane
     name: stable_manifold
 tags: [hide-input]
 ---

diff --git a/lectures/cass_koopmans_2.md b/lectures/cass_koopmans_2.md
@@ -68,7 +68,7 @@ Let's start with some standard imports:
 ```{code-cell} ipython
 import matplotlib.pyplot as plt
 plt.rcParams["figure.figsize"] = (11, 5)  #set default figure size
-from numba import njit, float64
+from numba import jit, float64
 from numba.experimental import jitclass
 import numpy as np
 ```
@@ -751,7 +751,7 @@ class PlanningProblem():
 ```
 
 ```{code-cell} python3
-@njit
+@jit
 def shooting(pp, c0, k0, T=10):
     '''
     Given the initial condition of capital k0 and an initial guess
@@ -779,7 +779,7 @@ def shooting(pp, c0, k0, T=10):
 ```
 
 ```{code-cell} python3
-@njit
+@jit
 def bisection(pp, c0, k0, T=10, tol=1e-4, max_iter=500, k_ter=0, verbose=True):
 
     # initial boundaries for guess c0
@@ -828,7 +828,7 @@ The above code from this lecture {doc}`Cass-Koopmans Planning Model <cass_koopma
 Now  we're ready to bring in Python code that we require to compute additional objects that appear in a competitive equilibrium.
 
 ```{code-cell} python3
-@njit
+@jit
 def q(pp, c_path):
     # Here we choose numeraire to be u'(c_0) -- this is q^(t_0)_t
     T = len(c_path) - 1
@@ -838,12 +838,12 @@ def q(pp, c_path):
         q_path[t] = pp.β ** t * pp.u_prime(c_path[t])
     return q_path
 
-@njit
+@jit
 def w(pp, k_path):
     w_path = pp.f(k_path) - k_path * pp.f_prime(k_path)
     return w_path
 
-@njit
+@jit
 def η(pp, k_path):
     η_path = pp.f_prime(k_path)
     return η_path
@@ -953,7 +953,7 @@ years, and define a new function for $r$, then plot both.
 We begin by continuing to assume that  $t_0=0$ and plot things for different maturities $t=T$, with $K_0$ below the steady state
 
 ```{code-cell} python3
-@njit
+@jit
 def q_generic(pp, t0, c_path):
     # simplify notations
     β = pp.β
@@ -966,7 +966,7 @@ def q_generic(pp, t0, c_path):
         q_path[t-t0] = β ** (t-t0) * u_prime(c_path[t]) / u_prime(c_path[t0])
     return q_path
 
-@njit
+@jit
 def r(pp, t0, q_path):
     '''Yield to maturity'''
     r_path = - np.log(q_path[1:]) / np.arange(1, len(q_path))

diff --git a/lectures/coleman_policy_iter.md b/lectures/coleman_policy_iter.md
@@ -63,7 +63,7 @@ Let's start with some imports:
 import matplotlib.pyplot as plt
 import numpy as np
 from quantecon.optimize import brentq
-from numba import njit
+from numba import jit
 ```
 
 ## The Euler Equation
@@ -286,7 +286,7 @@ u'(c) - \beta \int (u' \circ \sigma) (f(y - c) z ) f'(y - c) z \phi(dz)
 ```
 
 ```{code-cell} ipython
-@njit
+@jit
 def euler_diff(c, σ, y, og):
     """
     Set up a function such that the root with respect to c,
@@ -314,7 +314,7 @@ state $y$ and $σ$, the current guess of the policy.
 Here's the operator $K$, that implements the root-finding step.
 
 ```{code-cell} ipython3
-@njit
+@jit
 def K(σ, og):
     """
     The Coleman-Reffett operator

diff --git a/lectures/egm_policy_iter.md b/lectures/egm_policy_iter.md
@@ -44,7 +44,7 @@ Let's start with some standard imports:
 ```{code-cell} ipython
 import matplotlib.pyplot as plt
 import numpy as np
-from numba import njit
+from numba import jit
 ```
 
 ## Key Idea
@@ -161,7 +161,7 @@ We reuse the `OptimalGrowthModel` class
 Here's an implementation of $K$ using EGM as described above.
 
 ```{code-cell} python3
-@njit
+@jit
 def K(σ_array, og):
     """
     The Coleman-Reffett operator using EGM

diff --git a/lectures/eig_circulant.md b/lectures/eig_circulant.md
@@ -31,7 +31,7 @@ We begin by importing some Python packages
 
 ```{code-cell} ipython3
 import numpy as np
-from numba import njit
+from numba import jit
 import matplotlib.pyplot as plt
 ```
 
@@ -66,7 +66,7 @@ first column needs to be specified.
 Let's write some Python code to generate a circulant matrix.
 
 ```{code-cell} ipython3
-@njit
+@jit
 def construct_cirlulant(row):
 
     N = row.size
@@ -200,7 +200,7 @@ $$
 Let's write some Python code to illustrate these ideas.
 
 ```{code-cell} ipython3
-@njit
+@jit
 def construct_P(N):
 
     P = np.zeros((N, N))