From e6699ee140d2d5aff2c197399c408f54ce9959e2 Mon Sep 17 00:00:00 2001 From: Smit-create Date: Fri, 30 Dec 2022 16:12:24 +0530 Subject: [PATCH 01/11] add jax implementation for wealth dynamics time series --- lectures/wealth_dynamics.md | 144 +++++++++++++++++++++++++++++++++++- 1 file changed, 142 insertions(+), 2 deletions(-) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index c2a9954b4..e4cb76fc3 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -78,6 +78,9 @@ import numpy as np import quantecon as qe from numba import njit, float64, prange from numba.experimental import jitclass +import jax +import jax.numpy as jnp +from collections import namedtuple ``` ## Lorenz Curves and the Gini Coefficient @@ -250,7 +253,144 @@ their wealth. We are using something akin to a fixed savings rate model, while acknowledging that low wealth households tend to save very little. -## Implementation +## Implementation using JAX + +Let's define a Model to represent the wealth dynamics. + +```{code-cell} ipython3 +# NamedTuple Model +Model = namedtuple("Model", ("w_hat", "s_0", "c_y", "μ_y", + "σ_y", "c_r", "μ_r", "σ_r", "a", + "b", "σ_z", "z_mean", "z_var", "y_mean")) +``` + +Here's a function to create the Model with the given parameters + +```{code-cell} ipython3 +def create_wealth_model(w_hat=1.0, + s_0=0.75, + c_y=1.0, + μ_y=1.0, + σ_y=0.2, + c_r=0.05, + μ_r=0.1, + σ_r=0.5, + a=0.5, + b=0.0, + σ_z=0.1): + """ + Create a wealth model with given parameters and return + and instance of NamedTuple Model. + """ + z_mean = b / (1 - a) + z_var = σ_z**2 / (1 - a**2) + exp_z_mean = np.exp(z_mean + z_var / 2) + R_mean = c_r * exp_z_mean + np.exp(μ_r + σ_r**2 / 2) + y_mean = c_y * exp_z_mean + np.exp(μ_y + σ_y**2 / 2) + # Test a stability condition that ensures wealth does not diverge + # to infinity. + α = R_mean * s_0 + if α >= 1: + raise ValueError("Stability condition failed.") + return Model(w_hat=w_hat, s_0=s_0, c_y=c_y, μ_y=μ_y, + σ_y=σ_y, c_r=c_r, μ_r=μ_r, σ_r=σ_r, a=a, + b=b, σ_z=σ_z, z_mean=z_mean, z_var=z_var, y_mean=y_mean) +``` + +The following function updates one period with the given current wealth and persistent state. + +```{code-cell} ipython3 +def update_states_jax(arrays, wdy, size, rand_key): + """ + Update one period, given current wealth w and persistent + state z. They are stored in the form of tuples under the arrays argument + """ + # Unpack w and z + w, z = arrays + + rand_key, *subkey = jax.random.split(rand_key, 3) + zp = wdy.a * z + wdy.b + wdy.σ_z * jax.random.normal(rand_key, shape=size) + + # Update wealth + y = wdy.c_y * jnp.exp(zp) + jnp.exp(wdy.μ_y + wdy.σ_y * jax.random.normal(subkey[0], shape=size)) + wp = y + + R = wdy.c_r * jnp.exp(zp) + jnp.exp(wdy.μ_r + wdy.σ_r * jax.random.normal(subkey[1], shape=size)) + wp += (w >= wdy.w_hat) * R * wdy.s_0 * w + return wp, zp + +# Create the jit function +update_states_jax = jax.jit(update_states_jax, static_argnums=(2,)) +``` + +Here’s function to simulate the time series of wealth for in individual households. + +```{code-cell} ipython3 +def wealth_time_series_jax(w_0, n, wdy, size, rand_seed=1): + """ + Generate a single time series of length n for wealth given + initial value w_0. + + The initial persistent state z_0 for each household is drawn from + the stationary distribution of the AR(1) process. + + * wdy: NamedTuple Model + * w_0: scalar/vector + * n: int + * size: size/shape of the w_0 + * rand_seed: int (Used to generate PRNG key) + """ + rand_key = jax.random.PRNGKey(rand_seed) + rand_key, *subkey = jax.random.split(rand_key, n) + + w_0 = jax.device_put(w_0).reshape(size) + z_init = wdy.z_mean + jnp.sqrt(wdy.z_var) * jax.random.normal(rand_key, shape=size) + arrays = w_0, z_init + rand_sub_keys = jnp.array(subkey) + + w_final = jnp.array([w_0]) + + # Define the function for each update + def update_w_z(arrays, rand_sub_key): + wp, zp = update_states_jax(arrays, wdy, size, rand_sub_key) + return (wp, zp), wp + + arrays_last, w_values = jax.lax.scan(update_w_z, arrays, rand_sub_keys) + return jnp.concatenate((w_final, w_values)) + +# Create the jit function +wealth_time_series_jax = jax.jit(wealth_time_series_jax, static_argnums=(1,3,)) +``` + +Let's try simulating the model at different parameter values and investigate the implications for the wealth distribution. + +```{code-cell} ipython3 +wdy = create_wealth_model() # default model +ts_length = 200 +size = (1,) +``` + +```{code-cell} ipython3 +%%time + +w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size) +``` + +Running the above function again will be even faster because of JAX's JIT. + +```{code-cell} ipython3 +%%time + +# 2nd time is expected to be very fast because of JIT +w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size) +``` + +```{code-cell} ipython3 +fig, ax = plt.subplots() +ax.plot(w_jax_result) +plt.show() +``` +## Implementation using Numba Here's some type information to help Numba. @@ -648,4 +788,4 @@ plt.show() ``` ```{solution-end} -``` \ No newline at end of file +``` From 9e1657eb49958a3629a304795ca3a2f5aebaae56 Mon Sep 17 00:00:00 2001 From: Smit-create Date: Fri, 30 Dec 2022 16:59:56 +0530 Subject: [PATCH 02/11] Add update cross section in jax --- lectures/wealth_dynamics.md | 122 +++++++++++++++++++++++++++++++++++- 1 file changed, 121 insertions(+), 1 deletion(-) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index e4cb76fc3..b76153cdc 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -323,7 +323,7 @@ def update_states_jax(arrays, wdy, size, rand_key): update_states_jax = jax.jit(update_states_jax, static_argnums=(2,)) ``` -Here’s function to simulate the time series of wealth for in individual households. +Here’s function to simulate the time series of wealth for individual households. ```{code-cell} ipython3 def wealth_time_series_jax(w_0, n, wdy, size, rand_seed=1): @@ -390,6 +390,30 @@ fig, ax = plt.subplots() ax.plot(w_jax_result) plt.show() ``` + +Now here’s function to simulate a cross section of households forward in time. + +```{code-cell} ipython3 +def update_cross_section_jax(w_distribution, shift_length, wdy, size, rand_seed=2): + """ + Shifts a cross-section of household forward in time + + * wdy: NamedTuple Model + * w_distribution: array_like, represents current cross-section + + Takes a current distribution of wealth values as w_distribution + and updates each w_t in w_distribution to w_{t+j}, where + j = shift_length. + + Returns the new distribution. + """ + new_dist = wealth_time_series_jax(w_distribution, shift_length, wdy, size, rand_seed) + new_distribution = new_dist[-1, :] + return new_distribution + +# Create the jit function +update_cross_section_jax = jax.jit(update_cross_section_jax, static_argnums=(1,3,)) +``` ## Implementation using Numba Here's some type information to help Numba. @@ -558,7 +582,26 @@ wdy = WealthDynamics() ts_length = 200 w = wealth_time_series(wdy, wdy.y_mean, ts_length) +``` +```{code-cell} ipython3 +%%time + +w = wealth_time_series(wdy, wdy.y_mean, ts_length) +``` + +```{code-cell} ipython3 +%%time + +# Check the time for 2nd execution +w = wealth_time_series(wdy, wdy.y_mean, ts_length) +``` + +Notice the time difference between the `wealth_time_series` and `wealth_time_series_jax` + + + +```{code-cell} ipython3 fig, ax = plt.subplots() ax.plot(w) plt.show() @@ -575,7 +618,24 @@ Let's look at how inequality varies with returns on financial assets. The next function generates a cross section and then computes the Lorenz curve and Gini coefficient. +Let's first write the function that uses the jax implementation and for the numba. + ```{code-cell} ipython3 +# Uses jax +def generate_lorenz_and_gini_jax(wdy, num_households=100_000, T=500): + """ + Generate the Lorenz curve data and gini coefficient corresponding to a + WealthDynamics mode by simulating num_households forward to time T. + """ + size = (num_households, ) + ψ_0 = jnp.full(size, wdy.y_mean) + ψ_star = update_cross_section_jax(ψ_0, T, wdy, size) + ψ_star = jax.device_get(ψ_star) # get back to numpy form + return qe.gini_coefficient(ψ_star), qe.lorenz_curve(ψ_star) +``` + +```{code-cell} ipython3 +# Uses numba def generate_lorenz_and_gini(wdy, num_households=100_000, T=500): """ Generate the Lorenz curve data and gini coefficient corresponding to a @@ -599,6 +659,28 @@ This is unavoidable because we are executing a CPU intensive task. In fact the code, which is JIT compiled and parallelized, runs extremely fast relative to the number of computations. ```{code-cell} ipython3 +%%time + +fig, ax = plt.subplots() +μ_r_vals = (0.0, 0.025, 0.05) +gini_vals = [] + +for μ_r in μ_r_vals: + wdy = create_wealth_model(μ_r=μ_r) + gv, (f_vals, l_vals) = generate_lorenz_and_gini_jax(wdy) + ax.plot(f_vals, l_vals, label=f'$\psi^*$ at $\mu_r = {μ_r:0.2}$') + gini_vals.append(gv) + +ax.plot(f_vals, f_vals, label='equality') +ax.legend(loc="upper left") +plt.show() +``` + +Now let's try to run the same code snippet but using the numba version. + +```{code-cell} ipython3 +%%time + fig, ax = plt.subplots() μ_r_vals = (0.0, 0.025, 0.05) gini_vals = [] @@ -649,7 +731,31 @@ rise. Let's finish this section by investigating what happens when we change the volatility term $\sigma_r$ in financial returns. +Firstly, using jax, we have + ```{code-cell} ipython3 +%%time + +fig, ax = plt.subplots() +σ_r_vals = (0.35, 0.45, 0.52) +gini_vals = [] + +for σ_r in σ_r_vals: + wdy = create_wealth_model(σ_r=σ_r) + gv, (f_vals, l_vals) = generate_lorenz_and_gini_jax(wdy) + ax.plot(f_vals, l_vals, label=f'$\psi^*$ at $\sigma_r = {σ_r:0.2}$') + gini_vals.append(gv) + +ax.plot(f_vals, f_vals, label='equality') +ax.legend(loc="upper left") +plt.show() +``` + +Using numba we get, + +```{code-cell} ipython3 +%%time + fig, ax = plt.subplots() σ_r_vals = (0.35, 0.45, 0.52) gini_vals = [] @@ -764,6 +870,20 @@ z_0 = wdy.z_mean ``` First let's generate the distribution: +Using the jax implementation + +```{code-cell} ipython3 +num_households = 250_000 +T = 500 # how far to shift forward in time +size = (num_households, ) + +wdy = create_wealth_model() +ψ_0 = jnp.full(size, wdy.y_mean) +ψ_star = update_cross_section_jax(ψ_0, T, wdy, size) +ψ_star = jax.device_get(ψ_star) # get back numpy form +``` + +Using the numba implementation ```{code-cell} ipython3 num_households = 250_000 From a23e692a34a3510a7fa9daaed19010d984ec652a Mon Sep 17 00:00:00 2001 From: Smit-create Date: Sat, 31 Dec 2022 10:48:51 +0530 Subject: [PATCH 03/11] fix an error --- lectures/wealth_dynamics.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index b76153cdc..4bb4013a3 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -581,7 +581,6 @@ Let's look at the wealth dynamics of an individual household. wdy = WealthDynamics() ts_length = 200 -w = wealth_time_series(wdy, wdy.y_mean, ts_length) ``` ```{code-cell} ipython3 @@ -618,7 +617,7 @@ Let's look at how inequality varies with returns on financial assets. The next function generates a cross section and then computes the Lorenz curve and Gini coefficient. -Let's first write the function that uses the jax implementation and for the numba. +Let's first write the function that uses the jax implementation and then for the numba. ```{code-cell} ipython3 # Uses jax @@ -751,7 +750,7 @@ ax.legend(loc="upper left") plt.show() ``` -Using numba we get, +Using numba, we get, ```{code-cell} ipython3 %%time @@ -888,6 +887,7 @@ Using the numba implementation ```{code-cell} ipython3 num_households = 250_000 T = 500 # how far to shift forward in time +wdy = WealthDynamics() ψ_0 = np.full(num_households, wdy.y_mean) z_0 = wdy.z_mean From 2cb8379af1b2d67063970eff99f051e539ac7f6e Mon Sep 17 00:00:00 2001 From: Smit-create Date: Sat, 31 Dec 2022 11:43:53 +0530 Subject: [PATCH 04/11] add a cell to check backend --- lectures/wealth_dynamics.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index 4bb4013a3..46a0aed27 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -83,6 +83,15 @@ import jax.numpy as jnp from collections import namedtuple ``` +Let's check the backend used by jax and the devices available + +```{code-cell} ipython3 +# Check if JAX is using GPU +print("jax backend: {}".format(jax.lib.xla_bridge.get_backend().platform)) +# Check the devices available for JAX +print(jax.devices()) +``` + ## Lorenz Curves and the Gini Coefficient Before we investigate wealth dynamics, we briefly review some measures of @@ -869,6 +878,7 @@ z_0 = wdy.z_mean ``` First let's generate the distribution: + Using the jax implementation ```{code-cell} ipython3 From ff9d4c6f17181d1dea24c3eb825a04ff05eef9b8 Mon Sep 17 00:00:00 2001 From: Smit-create Date: Wed, 4 Jan 2023 09:55:15 +0530 Subject: [PATCH 05/11] Apply suggestions --- lectures/wealth_dynamics.md | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index 46a0aed27..2f1167fdf 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -83,11 +83,11 @@ import jax.numpy as jnp from collections import namedtuple ``` -Let's check the backend used by jax and the devices available +Let's check the backend used by JAX and the devices available ```{code-cell} ipython3 # Check if JAX is using GPU -print("jax backend: {}".format(jax.lib.xla_bridge.get_backend().platform)) +print(f"JAX backend: {jax.devices()[0].platform}") # Check the devices available for JAX print(jax.devices()) ``` @@ -626,10 +626,10 @@ Let's look at how inequality varies with returns on financial assets. The next function generates a cross section and then computes the Lorenz curve and Gini coefficient. -Let's first write the function that uses the jax implementation and then for the numba. +Let's first write the function that uses the JAX implementation and then for the numba. ```{code-cell} ipython3 -# Uses jax +# Uses JAX def generate_lorenz_and_gini_jax(wdy, num_households=100_000, T=500): """ Generate the Lorenz curve data and gini coefficient corresponding to a @@ -739,7 +739,7 @@ rise. Let's finish this section by investigating what happens when we change the volatility term $\sigma_r$ in financial returns. -Firstly, using jax, we have +Firstly, using JAX, we have ```{code-cell} ipython3 %%time @@ -879,7 +879,7 @@ z_0 = wdy.z_mean First let's generate the distribution: -Using the jax implementation +Using the JAX implementation ```{code-cell} ipython3 num_households = 250_000 From 063367f19f6a1774930b3c5585010dae86be887b Mon Sep 17 00:00:00 2001 From: Smit-create Date: Wed, 4 Jan 2023 10:05:35 +0530 Subject: [PATCH 06/11] Add for-loop implementation --- lectures/wealth_dynamics.md | 71 +++++++++++++++++++++++++++++++++++-- 1 file changed, 69 insertions(+), 2 deletions(-) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index 2f1167fdf..29b0472c1 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -332,7 +332,72 @@ def update_states_jax(arrays, wdy, size, rand_key): update_states_jax = jax.jit(update_states_jax, static_argnums=(2,)) ``` -Here’s function to simulate the time series of wealth for individual households. +Here’s function to simulate the time series of wealth for individual households using `for` loop and JAX. + +```{code-cell} ipython3 +# Using JAX and for loop +def wealth_time_series_for_loop_jax(w_0, n, wdy, size, rand_seed=1): + """ + Generate a single time series of length n for wealth given + initial value w_0. + + * This implementation uses for loop. + + The initial persistent state z_0 for each household is drawn from + the stationary distribution of the AR(1) process. + + * wdy: NamedTuple Model + * w_0: scalar/vector + * n: int + * size: size/shape of the w_0 + * rand_seed: int (Used to generate PRNG key) + """ + rand_key = jax.random.PRNGKey(rand_seed) + rand_key, *subkey = jax.random.split(rand_key, n) + + w_0 = jax.device_put(w_0).reshape(size) + + z = wdy.z_mean + jnp.sqrt(wdy.z_var) * jax.random.normal(rand_key, shape=size) + w = [w_0] + for t in range(n-1): + w_, z = update_states_jax((w[t], z), wdy, size, subkey[t]) + w.append(w_) + return jnp.array(w) + +# Create the jit function +wealth_time_series_for_loop_jax = jax.jit(wealth_time_series_for_loop_jax, static_argnums=(1,3,)) +``` + +Let's try simulating the model at different parameter values and investigate the implications for the wealth distribution using the above function. + +```{code-cell} ipython3 +wdy = create_wealth_model() # default model +ts_length = 200 +size = (1,) +``` + +```{code-cell} ipython3 +%%time + +w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size) +``` + +Running the above function again will be even faster because of JAX's JIT. + +```{code-cell} ipython3 +%%time + +# 2nd time is expected to be very fast because of JIT +w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size) +``` + +```{code-cell} ipython3 +fig, ax = plt.subplots() +ax.plot(w_jax_result) +plt.show() +``` + +We can further try to optimize and speed up the compile time of the above function by replacing `for` loop with [`jax.lax.scan`](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html). ```{code-cell} ipython3 def wealth_time_series_jax(w_0, n, wdy, size, rand_seed=1): @@ -340,6 +405,8 @@ def wealth_time_series_jax(w_0, n, wdy, size, rand_seed=1): Generate a single time series of length n for wealth given initial value w_0. + * This implementation uses for jax.lax.scan + The initial persistent state z_0 for each household is drawn from the stationary distribution of the AR(1) process. @@ -371,7 +438,7 @@ def wealth_time_series_jax(w_0, n, wdy, size, rand_seed=1): wealth_time_series_jax = jax.jit(wealth_time_series_jax, static_argnums=(1,3,)) ``` -Let's try simulating the model at different parameter values and investigate the implications for the wealth distribution. +Let's try simulating the model at different parameter values and investigate the implications for the wealth distribution and also observe the difference in time between `wealth_time_series_jax` and `wealth_time_series_for_loop_jax`. ```{code-cell} ipython3 wdy = create_wealth_model() # default model From d9423834e8ec833d5e517ed7702f99cd7271bf46 Mon Sep 17 00:00:00 2001 From: Smit-create Date: Wed, 4 Jan 2023 12:22:32 +0530 Subject: [PATCH 07/11] Hide numba code --- lectures/wealth_dynamics.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index 29b0472c1..520b352dd 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -495,6 +495,9 @@ update_cross_section_jax = jax.jit(update_cross_section_jax, static_argnums=(1,3 Here's some type information to help Numba. ```{code-cell} ipython3 +--- +tags: [hide-input] +--- wealth_dynamics_data = [ ('w_hat', float64), # savings parameter ('s_0', float64), # savings parameter @@ -518,6 +521,9 @@ Here's a class that stores instance data and implements methods that update the aggregate state and household wealth. ```{code-cell} ipython3 +--- +tags: [hide-input] +--- @jitclass(wealth_dynamics_data) class WealthDynamics: @@ -613,6 +619,9 @@ Now here's function to simulate a cross section of households forward in time. Note the use of parallelization to speed up computation. ```{code-cell} ipython3 +--- +tags: [hide-input] +--- @njit(parallel=True) def update_cross_section(wdy, w_distribution, shift_length=500): """ @@ -709,7 +718,12 @@ def generate_lorenz_and_gini_jax(wdy, num_households=100_000, T=500): return qe.gini_coefficient(ψ_star), qe.lorenz_curve(ψ_star) ``` +The following function uses the numba implementation + ```{code-cell} ipython3 +--- +tags: [hide-input] +--- # Uses numba def generate_lorenz_and_gini(wdy, num_households=100_000, T=500): """ @@ -754,6 +768,9 @@ plt.show() Now let's try to run the same code snippet but using the numba version. ```{code-cell} ipython3 +--- +tags: [hide-input] +--- %%time fig, ax = plt.subplots() @@ -829,6 +846,9 @@ plt.show() Using numba, we get, ```{code-cell} ipython3 +--- +tags: [hide-input] +--- %%time fig, ax = plt.subplots() From e43c237ef48acfc6a1f251143e05f7271267382c Mon Sep 17 00:00:00 2001 From: Smit-create Date: Tue, 10 Jan 2023 11:34:13 +0530 Subject: [PATCH 08/11] Hide all numba code and fix typos --- lectures/wealth_dynamics.md | 37 ++++++++++++++++++++----------------- 1 file changed, 20 insertions(+), 17 deletions(-) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index 520b352dd..71d3d57e1 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -264,7 +264,7 @@ acknowledging that low wealth households tend to save very little. ## Implementation using JAX -Let's define a Model to represent the wealth dynamics. +Let's define a model to represent the wealth dynamics. ```{code-cell} ipython3 # NamedTuple Model @@ -332,7 +332,7 @@ def update_states_jax(arrays, wdy, size, rand_key): update_states_jax = jax.jit(update_states_jax, static_argnums=(2,)) ``` -Here’s function to simulate the time series of wealth for individual households using `for` loop and JAX. +Here’s function to simulate the time series of wealth for individual households using a `for` loop and JAX. ```{code-cell} ipython3 # Using JAX and for loop @@ -341,7 +341,7 @@ def wealth_time_series_for_loop_jax(w_0, n, wdy, size, rand_seed=1): Generate a single time series of length n for wealth given initial value w_0. - * This implementation uses for loop. + * This implementation uses a for loop. The initial persistent state z_0 for each household is drawn from the stationary distribution of the AR(1) process. @@ -379,7 +379,7 @@ size = (1,) ```{code-cell} ipython3 %%time -w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size) +w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() ``` Running the above function again will be even faster because of JAX's JIT. @@ -388,7 +388,7 @@ Running the above function again will be even faster because of JAX's JIT. %%time # 2nd time is expected to be very fast because of JIT -w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size) +w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() ``` ```{code-cell} ipython3 @@ -449,7 +449,7 @@ size = (1,) ```{code-cell} ipython3 %%time -w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size) +w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() ``` Running the above function again will be even faster because of JAX's JIT. @@ -458,7 +458,7 @@ Running the above function again will be even faster because of JAX's JIT. %%time # 2nd time is expected to be very fast because of JIT -w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size) +w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() ``` ```{code-cell} ipython3 @@ -591,6 +591,9 @@ class WealthDynamics: Here's function to simulate the time series of wealth for in individual households. ```{code-cell} ipython3 +--- +tags: [hide-input] +--- @njit def wealth_time_series(wdy, w_0, n): """ @@ -660,21 +663,30 @@ the implications for the wealth distribution. ### Time Series -Let's look at the wealth dynamics of an individual household. +Let's look at the wealth dynamics of an individual household using numba. ```{code-cell} ipython3 +--- +tags: [hide-input] +--- wdy = WealthDynamics() ts_length = 200 ``` ```{code-cell} ipython3 +--- +tags: [hide-input] +--- %%time w = wealth_time_series(wdy, wdy.y_mean, ts_length) ``` ```{code-cell} ipython3 +--- +tags: [hide-input] +--- %%time # Check the time for 2nd execution @@ -794,15 +806,6 @@ We will look at this again via the Gini coefficient immediately below, but first consider the following image of our system resources when the code above is executing: -(htop_again)= -```{figure} /_static/lecture_specific/wealth_dynamics/htop_again.png -:scale: 80 -``` - -Notice how effectively Numba has implemented multithreading for this routine: -all 8 CPUs on our workstation are running at maximum capacity (even though -four of them are virtual). - Since the code is both efficiently JIT compiled and fully parallelized, it's close to impossible to make this sequence of tasks run faster without changing hardware. From 975be525b251483934c8feb198d45c38914499c3 Mon Sep 17 00:00:00 2001 From: Smit-create Date: Fri, 13 Jan 2023 20:48:03 +0530 Subject: [PATCH 09/11] Remove hidden tags and use glue --- lectures/wealth_dynamics.md | 60 +++++++++++-------------------------- 1 file changed, 18 insertions(+), 42 deletions(-) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index 71d3d57e1..d6d92edbc 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -81,6 +81,7 @@ from numba.experimental import jitclass import jax import jax.numpy as jnp from collections import namedtuple +from myst_nb import glue ``` Let's check the backend used by JAX and the devices available @@ -377,18 +378,18 @@ size = (1,) ``` ```{code-cell} ipython3 -%%time - +qe.tic() w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() +qe.toc() ``` Running the above function again will be even faster because of JAX's JIT. ```{code-cell} ipython3 -%%time - +qe.tic() # 2nd time is expected to be very fast because of JIT w_jax_result = wealth_time_series_for_loop_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() +qe.toc() ``` ```{code-cell} ipython3 @@ -448,17 +449,18 @@ size = (1,) ```{code-cell} ipython3 %%time - +qe.tic() w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() +glue("wealth_time_series_jax_time_1", qe.toc()) ``` Running the above function again will be even faster because of JAX's JIT. ```{code-cell} ipython3 -%%time - +qe.tic() # 2nd time is expected to be very fast because of JIT w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() +glue("wealth_time_series_jax_time_2", qe.toc()) ``` ```{code-cell} ipython3 @@ -495,9 +497,7 @@ update_cross_section_jax = jax.jit(update_cross_section_jax, static_argnums=(1,3 Here's some type information to help Numba. ```{code-cell} ipython3 ---- -tags: [hide-input] ---- + wealth_dynamics_data = [ ('w_hat', float64), # savings parameter ('s_0', float64), # savings parameter @@ -521,9 +521,7 @@ Here's a class that stores instance data and implements methods that update the aggregate state and household wealth. ```{code-cell} ipython3 ---- -tags: [hide-input] ---- + @jitclass(wealth_dynamics_data) class WealthDynamics: @@ -591,9 +589,7 @@ class WealthDynamics: Here's function to simulate the time series of wealth for in individual households. ```{code-cell} ipython3 ---- -tags: [hide-input] ---- + @njit def wealth_time_series(wdy, w_0, n): """ @@ -622,9 +618,7 @@ Now here's function to simulate a cross section of households forward in time. Note the use of parallelization to speed up computation. ```{code-cell} ipython3 ---- -tags: [hide-input] ---- + @njit(parallel=True) def update_cross_section(wdy, w_distribution, shift_length=500): """ @@ -666,34 +660,25 @@ the implications for the wealth distribution. Let's look at the wealth dynamics of an individual household using numba. ```{code-cell} ipython3 ---- -tags: [hide-input] ---- wdy = WealthDynamics() ts_length = 200 ``` ```{code-cell} ipython3 ---- -tags: [hide-input] ---- -%%time - +qe.tic() w = wealth_time_series(wdy, wdy.y_mean, ts_length) +glue("wealth_time_series_time_1", qe.toc()) ``` ```{code-cell} ipython3 ---- -tags: [hide-input] ---- -%%time - +qe.tic() # Check the time for 2nd execution w = wealth_time_series(wdy, wdy.y_mean, ts_length) +glue("wealth_time_series_time_2", qe.toc()) ``` -Notice the time difference between the `wealth_time_series` and `wealth_time_series_jax` +Notice the time difference between the `wealth_time_series`: {glue:}`wealth_time_series_time_1` and `wealth_time_series_jax`: {glue:}`wealth_time_series_jax_time_1` @@ -733,9 +718,6 @@ def generate_lorenz_and_gini_jax(wdy, num_households=100_000, T=500): The following function uses the numba implementation ```{code-cell} ipython3 ---- -tags: [hide-input] ---- # Uses numba def generate_lorenz_and_gini(wdy, num_households=100_000, T=500): """ @@ -780,9 +762,6 @@ plt.show() Now let's try to run the same code snippet but using the numba version. ```{code-cell} ipython3 ---- -tags: [hide-input] ---- %%time fig, ax = plt.subplots() @@ -849,9 +828,6 @@ plt.show() Using numba, we get, ```{code-cell} ipython3 ---- -tags: [hide-input] ---- %%time fig, ax = plt.subplots() From 6b49755aac6cfc9a5ebe8586551ad0820f7721bc Mon Sep 17 00:00:00 2001 From: Smit-create Date: Mon, 16 Jan 2023 11:23:10 +0530 Subject: [PATCH 10/11] Add GPU warning --- lectures/wealth_dynamics.md | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index d6d92edbc..db874c0c0 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -19,6 +19,15 @@ kernelspec: # Wealth Distribution Dynamics +```{admonition} GPU Warning +:class: dropdown, warning +This lecture is built using [hardware](status:machine-details) that has access to a GPU and uses JAX for GPU programming. As a result, the lecture will be slower when running on a machine without a GPU. + +Free GPUs are available on Google Colab. To use this option, please click on the play icon top right, select Colab, and set the runtime environment to include a GPU. + +Alternatively, if you have your own GPU, you can follow the [instructions](https://github.com/google/jax#pip-installation-gpu-cuda) for installing JAX with GPU support. +``` + ```{contents} Contents :depth: 2 ``` @@ -30,6 +39,9 @@ In addition to what's in Anaconda, this lecture will need the following librarie tags: [hide-output] --- !pip install quantecon +# If your machine has CUDA support, please follow the guide in GPU Warning. +# Otherwise, run the line below: +!pip install --upgrade "jax[CPU]" ``` ## Overview @@ -448,7 +460,6 @@ size = (1,) ``` ```{code-cell} ipython3 -%%time qe.tic() w_jax_result = wealth_time_series_jax(wdy.y_mean, ts_length, wdy, size).block_until_ready() glue("wealth_time_series_jax_time_1", qe.toc()) From d4e56bae3498d45cb7afea13028744cd8a39921e Mon Sep 17 00:00:00 2001 From: Smit-create Date: Tue, 17 Jan 2023 09:31:23 +0530 Subject: [PATCH 11/11] Update GPU warning --- lectures/status.md | 7 +++++-- lectures/wealth_dynamics.md | 7 ++++--- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/lectures/status.md b/lectures/status.md index b8a415a8f..d37a0b808 100644 --- a/lectures/status.md +++ b/lectures/status.md @@ -16,5 +16,8 @@ This table contains the latest execution statistics. ```{nb-exec-table} ``` -These lectures are built on `linux` instances through `github actions` so are -executed using the following [hardware specifications](https://docs.github.com/en/actions/using-github-hosted-runners/about-github-hosted-runners#supported-runners-and-hardware-resources) \ No newline at end of file +(status:machine-details)= + +These lectures are built on `linux` instances through `github actions` and `amazon web services (aws)` to +enable access to a `gpu`. These lectures are built on a [p3.2xlarge](https://aws.amazon.com/ec2/instance-types/p3/) +that has access to `8 vcpu's`, a `V100 NVIDIA Tesla GPU`, and `61 Gb` of memory. diff --git a/lectures/wealth_dynamics.md b/lectures/wealth_dynamics.md index db874c0c0..6bf34e132 100644 --- a/lectures/wealth_dynamics.md +++ b/lectures/wealth_dynamics.md @@ -19,9 +19,10 @@ kernelspec: # Wealth Distribution Dynamics -```{admonition} GPU Warning -:class: dropdown, warning -This lecture is built using [hardware](status:machine-details) that has access to a GPU and uses JAX for GPU programming. As a result, the lecture will be slower when running on a machine without a GPU. +```{admonition} GPU in use +:class: warning + +This lecture is accelerated via [hardware](status:machine-details) that has access to a GPU and JAX for GPU programming. Free GPUs are available on Google Colab. To use this option, please click on the play icon top right, select Colab, and set the runtime environment to include a GPU.