You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
I was reading through the wealth_dynamics lecture and had a few small suggestions.
I’d be grateful if you could take a look when you have a chance.
Opening paragraph
It might be helpful to add one sentence explaining why this lecture is a natural fit for JAX: the model involves simulating large cross sections of households, which suits JAX's array-based programming model and accelerator support.
Wealth dynamics
In the savings rule explanation, w \geq \bar w should likely be w \geq \hat w, matching the threshold in the displayed savings function.
It may be helpful to clarify the timing convention: when updating from $w_t$ to $w_{t+1}$, the simulation draws $z_{t+1}$ and the period $t+1$ idiosyncratic shocks.
Numba / JAX implementation
The simulation horizon and aggregate-state indexing seem inconsistent across the individual path, Numba cross-section, and JAX cross-section implementations.
The model equation updates wealth from $w_t$ to $w_{t+1}$ using next-period shocks:
$$w_{t+1} = (1 + r_{t+1}) s(w_t) + y_{t+1}$$
The function generate_aggregate_state_sequence(..., length=T) returns T + 1 aggregate states. It initializes z[0] at the stationary mean, then simulates z[1], \ldots, z[T].
This convention is used in the individual time-series simulation, where the update uses z[t+1]. In other words, z[0] is treated as the initial aggregate state, and the wealth updates use the simulated next-period aggregate states.
The cross-section implementations seem to use a different convention. The Numba cross-section loop updates with z[t], so the first update uses z[0] and the last simulated aggregate state is not used. The JAX cross-section implementations loop over the full z_sequence, so they also use z[0] and appear to apply one more update than the Numba cross-section code.
It may be clearer to choose one timing convention and apply it consistently across the individual, Numba, and JAX simulations. For example, if z[0] is the initial aggregate state, then all wealth updates could use z_sequence[1:], corresponding to the aggregate states used when moving from $w_t$ to $w_{t+1}$.
In the first JAX implementation, the PRNG key is used before splitting and the resulting subkey is unused. It may be clearer to use the standard JAX pattern: split first, then draw with the subkey.
Exercises
In Exercises 2 and 3, ax.plot(x, y, label='equality') plots the last Lorenz curve again rather than the 45-degree equality line as shown below. This should likely be ax.plot(x, x, label='equality').
In Exercise 3, setting $\mu_r = -\sigma_r^2 / 2$ holds the mean of the lognormal idiosyncratic return component constant, not its variance. The parenthetical should likely say “mean” rather than “variance”.
In the exercise solutions, the same PRNG key is reused across parameter values. This is useful as a common-random-number comparison, but it may be worth stating explicitly so readers do not interpret it as accidental key reuse.
Hi @jstac
I was reading through the
wealth_dynamicslecture and had a few small suggestions.I’d be grateful if you could take a look when you have a chance.
Opening paragraph
Wealth dynamics
In the savings rule explanation,
w \geq \bar wshould likely bew \geq \hat w, matching the threshold in the displayed savings function.It may be helpful to clarify the timing convention: when updating from$w_t$ to $w_{t+1}$ , the simulation draws $z_{t+1}$ and the period $t+1$ idiosyncratic shocks.
Numba / JAX implementation
The simulation horizon and aggregate-state indexing seem inconsistent across the individual path, Numba cross-section, and JAX cross-section implementations.
The model equation updates wealth from$w_t$ to $w_{t+1}$ using next-period shocks:
The function
generate_aggregate_state_sequence(..., length=T)returnsT + 1aggregate states. It initializesz[0]at the stationary mean, then simulatesz[1], \ldots, z[T].This convention is used in the individual time-series simulation, where the update uses
z[t+1]. In other words,z[0]is treated as the initial aggregate state, and the wealth updates use the simulated next-period aggregate states.The cross-section implementations seem to use a different convention. The Numba cross-section loop updates with
z[t], so the first update usesz[0]and the last simulated aggregate state is not used. The JAX cross-section implementations loop over the fullz_sequence, so they also usez[0]and appear to apply one more update than the Numba cross-section code.It may be clearer to choose one timing convention and apply it consistently across the individual, Numba, and JAX simulations. For example, if$w_t$ to $w_{t+1}$ .
z[0]is the initial aggregate state, then all wealth updates could usez_sequence[1:], corresponding to the aggregate states used when moving fromIn the first JAX implementation, the PRNG key is used before splitting and the resulting
subkeyis unused. It may be clearer to use the standard JAX pattern: split first, then draw with the subkey.Exercises
ax.plot(x, y, label='equality')plots the last Lorenz curve again rather than the 45-degree equality line as shown below. This should likely beax.plot(x, x, label='equality').In Exercise 3, setting$\mu_r = -\sigma_r^2 / 2$ holds the mean of the lognormal idiosyncratic return component constant, not its variance. The parenthetical should likely say “mean” rather than “variance”.
In the exercise solutions, the same PRNG key is reused across parameter values. This is useful as a common-random-number comparison, but it may be worth stating explicitly so readers do not interpret it as accidental key reuse.
What do you think? Happy to submit a PR.
Best,
Longye