Tutorial: Two-period model¶
The purpose of this tutorial is to demonstrate the accuracy of the DCEGM numerical solution method by comparing it with an analytical solution for a simple two-period model.
We begin by presenting the theoretical model. Next, we demonstrate how to solve it numerically using dcegm. Finally, we derive its analytical solution confirming the equivalence of the two approaches.
Two-period model¶
We consider a simple consumption-retirement model with only two periods. The objective function in period 0 is given by
where \(\beta \in [0,1]\) is the discount factor, \(c_t\geq 0\) is the consumption in period \(t\) and \(d_t\in \{0,1\}\) specifies the discrete choice with \(d_t = 0\) indicating work and \(d_t=1\) indicating retirement at the end of period \(t\).
The utility function has the form
The parameters \(\rho \geq 0\) and \(\delta \geq 0\) are measures of risk aversion and disutility of work, respectively, while \(\epsilon \sim EV(0,1)\) is a choice-specific taste shock with extreme-value distribution.
In each period \(t\), the consumption \(c_t\) has to satisfy the budget constraint \(c_t \leq M_t\) with wealth
where \(R\) is the interest factor, \(W_t\) is the wage in period \(t\) and \(D_t\) is an exogenous process indicating long-term care dependence with cost \(K\). The wage \(W_t = W+\nu_t\) consists of an average component \(W\), which in the two-period model does not depend on \(t\) as only the wage in period 1 enters the period 2 budget constraint, and an income shock \(\nu_t \sim EV(0,1)\). For the exogenous probability of becoming long-term care dependent, we have the following transition probabilities \(\pi\)
where \(\pi(D_t=1\mid D_{t-1}=1)=1\) means that care dependence is absorbing (similar to the retirement decision).
import jax.numpy as jnp
import numpy as np
import dcegm
from scipy.special import roots_sh_legendre
from scipy.stats import norm
from typing import Tuple
import jax
jax.config.update("jax_enable_x64", True)
Solve model using DCEGM¶
We now demonstrate how to solve the two-period model outlined above using the dcegm package. First, we define the parameters of the model and store them in a dictionary. These parameters are accessed in user-defined functions through the params argument, which is a required input for all user functions.
params = {}
params["interest_rate"] = 0.02
params["ltc_cost"] = 5
params["max_wealth"] = 50
params["wage_avg"] = 8
params["income_shock_std"] = 1
params["income_shock_mean"] = 0
params["taste_shock_scale"] = 1
params["ltc_prob"] = 0.3
params["discount_factor"] = 0.95
params["rho"] = 0.9
params["delta"] = 1.5
Second, the user needs to specify a model_config dictionary containing at least the following information:
number of periods,
the discrete choices,
optional: dictionaries which contain the deterministic and the stochastic states
size of the exogenous wealth grid,
number of stochastic quadrature points
Additionally, the user can specify "model_specs", which are passed into all user-supplied functions. These parameters are considered constant and independent of the parametrization of the model.
# Exogenous process: Long-term care dependence (see above)
def prob_exog_ltc(ltc, params):
prob_ltc = params["ltc_prob"]
ltc_depenent = ltc == 1
prob_ltc = ltc_depenent + (1 - ltc_depenent) * prob_ltc
return jnp.array([1 - prob_ltc, prob_ltc])
model_specs = {
"choices": np.arange(2),
}
model_config = {
"n_periods": 2,
"choices": np.arange(2),
"stochastic_states": {
"ltc": [0, 1],
},
"continuous_states": {
"assets_end_of_period": np.linspace(0, 50, 100),
},
"n_quad_points": 5,
}
We will describe the relevant keywords for each user-defined function in the following sections.
Beside the params and options dictionaries, the main function of the dcegm package solve_dcegm requires the following inputs:
utility_functions,
budget_constraint,
solve_final_period,
state_space_functions
In the following, we describe the format of the required functions. Note that they must be JAX-compatible, i.e., pure functions that avoid side effects, such as conditional statements. For more details, refer to the Pure Functions section in the JAX documentation.
Utility functions¶
First, we define the utility, the marginal utility and inverse marginal utility. These functions are stored in a dictionary called utility_functions.
def flow_util(consumption, choice, params):
rho = params["rho"]
delta = params["delta"]
u = consumption ** (1 - rho) / (1 - rho) - delta * (1 - choice)
return u
def marginal_utility(consumption, params):
rho = params["rho"]
u_prime = consumption ** (-rho)
return u_prime
def inverse_marginal_utility(marginal_utility, params):
rho = params["rho"]
return marginal_utility ** (-1 / rho)
utility_functions = {
"utility": flow_util,
"inverse_marginal_utility": inverse_marginal_utility,
"marginal_utility": marginal_utility,
}
State space functions¶
The state_specific_choice_set function specifies the set of feasible choices for each state. It accounts for the fact that retirement is absorbing - that is, if \(d_{t-1}=1\), then it must also hold that \(d_{t}=1\).
def state_specific_choice_set(
lagged_choice: np.ndarray,
model_specs: np.ndarray,
) -> np.ndarray:
"""Select state-specific choice set.
Args:
state (np.ndarray): Array of shape (n_state_variables,) defining the agent's
state. In Ishkakov, an agent's state is defined by her (i) age (i.e., the
current period) and (ii) her lagged labor market choice.
Hence n_state_variables = 2.
state_space (np.ndarray): Collection of all possible states of shape
(n_periods * n_choices, n_choices).
indexer (np.ndarray): Indexer object that maps states to indexes.
Shape (n_periods, n_choices).
Returns:
choice_set (np.ndarray): The agent's (restricted) choice set in the given
state of shape (n_admissible_choices,).
"""
# Once the agent choses retirement, she can only choose retirement thereafter.
# Hence, retirement is an absorbing state.
if lagged_choice == 1:
choice_set = np.array([1])
else:
choice_set = model_specs["choices"]
return choice_set
As an example, consider a where the agent chooses to work in the first period, i.e., \(d_0 = 0\). Consequently, in the second period, both choices - \(d_1 = 0\) (continue working), \(d_1 = 1\) (retire) - are possible.
choice_set_worker = state_specific_choice_set(lagged_choice=0, model_specs=model_specs)
choice_set_worker
array([0, 1])
Let’s now consider a case where \(d_0 = 1\). In the second period, now only the choice \(d_1 = 1\) is possible.
choice_set_retiree = state_specific_choice_set(lagged_choice=1, model_specs=model_specs)
choice_set_retiree
array([1])
The state_specific_choice_set function needs to be placed in a dictionary called state_space_functions before being passed to the main function solve_dcegm.
state_space_functions = {
"state_specific_choice_set": state_specific_choice_set,
}
Budget function and transition function¶
Next, we define the budget and transition functions, which take the params dictionary as input.
def budget_constraint(
lagged_choice,
ltc,
asset_end_of_previous_period,
income_shock_previous_period,
params,
):
interest_factor = 1 + params["interest_rate"]
health_costs = params["ltc_cost"]
wage = params["wage_avg"]
resource = (
interest_factor * asset_end_of_previous_period
+ (wage + income_shock_previous_period) * (1 - lagged_choice)
- ltc * health_costs
)
return jnp.maximum(resource, 0.5)
# Already defined above; repeated here for completeness
def prob_exog_ltc(ltc, params): # noqa: F811
prob_ltc = params["ltc_prob"]
ltc_depenent = ltc == 1
prob_ltc = ltc_depenent + (1 - ltc_depenent) * prob_ltc
return jnp.array([1 - prob_ltc, prob_ltc])
stochastic_states_transitions = {
"ltc": prob_exog_ltc,
}
Solve final period function¶
Lastly, we need a function that computes the solution in the final period. While this function can be imported directly from the package, we also present it here for clarity.
def final_period_utility(wealth: float, choice: int, params) -> Tuple[float, float]:
return flow_util(wealth, choice, params)
def marginal_final(wealth, choice):
return marginal_utility(wealth, params)
utility_functions_final_period = {
"utility": final_period_utility,
"marginal_utility": marginal_final,
}
Solve function¶
Once all input arguments have the required form (see above), they can be passed to the function solve_dcegm. This function returns two multi-dimensional arrays:
policy (np.ndarray): Multi-dimensional np.ndarray storing the choice-specific policy function; of shape [n_states, n_discrete_choices, 2, 1.1 * n_grid_wealth]. Position \([.., 0, :]\) contains the endogenous grid over wealth M, and \([.., 1, :]\) stores the corresponding value of the policy function c(M, d), for each state and each discrete choice.
value (np.ndarray): Multi-dimensional np.ndarray storing the choice-specific value functions; of shape [n_states, n_discrete_choices, 2, 1.1 * n_grid_wealth]. Position \([.., 0, :]\) contains the endogenous grid over wealth M, and \([.., 1, :]\) stores the corresponding value of the value function v(M, d), for each state and each discrete choice.
model = dcegm.setup_model(
model_config=model_config,
model_specs=model_specs,
utility_functions=utility_functions,
utility_functions_final_period=utility_functions_final_period,
state_space_functions=state_space_functions,
stochastic_states_transitions=stochastic_states_transitions,
budget_constraint=budget_constraint,
)
Update function for state space not given. Assume states only change with an increase of the period and lagged choice.
Sparsity condition not provided. Assume all states are valid.
Starting state space creation
State space created.
Starting state-choice space creation and child state mapping.
State, state-choice and child state mapping created.
Start creating batches for the model.
Model setup complete.
solved_model = model.solve(params)
value_function = solved_model.value
policy_function = solved_model.policy
Analytical solution of the model¶
The solution of the given problem can be derived analytically using backward induction.
Period 1:¶
The choice problem in period 1 can be expressed through the following Bellman equation
Since this is the final period of the model and there is no bequest, agents consume the entire budget, i.e., \(c_1 = M_1\). Therefore, the choice-specific value function for a given wealth level \(M_1\) and discrete choice \(d_1\) is given by
Period 0:¶
Analogous to period 1, the choice problem in period 0 can be expressed through the Bellman equation
Here, the choice-specific value function for a given wealth level \(M_0\) and choice \(d_0\) is defined by
where \(EV_1(M_1(\nu_1,D_1))\) is the expected value function for a given realization of the income shock \(\nu_1\) and exogenous process \(D_1\), i.e., it is the expected maximum of the different choice-specific value functions in the second period.
The extreme value distribution takes the following closed formulas for the expected value function and choice probabilities:
Now, the problem can be solved using the Euler equation given by (see Iskhakov et al., 2017, Appendix A, Lemma 1)
Policy functions¶
Using the fact that the marginal utility is given by \(u^\prime(c_t) = c_t^{-\rho}\) we obtain the consumption policy for period 0:
where we simplified the last equation using the second-period budget constraint \(M_1 = c_1\). Note that this policy function is implicit, as \(M_1\) depends on \(c_0\). More specifically, we have
The labor supply in period 0 \(d_0 \in \{0,1\}\) is the maximizer of \(v_0(M_0,d_0)+\epsilon_0(d_0)\). Hence, \(d_0 = 0\) if \(v_0(M_0,0)+\epsilon_0(0)\geq v_0(M_0,1)+\epsilon_0(1)\) and \(d_0 = 1\) otherwise.
Given \(c_0\) and \(d_0\), the consumption \(c_1\) in period 1, which is equal to the wealth \(M_1\) can be calculated directly using the budget constraint. The labor supply decision \(d_1\) in period 1 can be determined analogously to period 1 as the maximizer of \(v_1(M_1,d_1)+\epsilon_1(d_1)\).
The budget constraint, wage, transition probability, choice probabilities and the right-hand side of the Euler equation can be implemented as follows:
def budget(
lagged_resources, lagged_consumption, lagged_choice, wage, health, params_dict
):
end_of_previous_period_assets = lagged_resources - lagged_consumption
interest_factor = 1 + params_dict["interest_rate"]
health_costs = params_dict["ltc_cost"]
resources = (
interest_factor * end_of_previous_period_assets
+ wage * (1 - lagged_choice)
- health * health_costs
).clip(min=0.5)
return resources
def wage(nu, params_dict):
wage = params_dict["wage_avg"] + nu
return wage
def prob_long_term_care_patient(params_dict, lag_health, health):
p = params_dict["ltc_prob"]
if (lag_health == 0) and (health == 1):
pi = p
elif (lag_health == 0) and (health == 0):
pi = 1 - p
elif (lag_health == 1) and (health == 0):
pi = 0
elif (lag_health == 1) and (health == 1):
pi = 1
else:
raise ValueError("Health state not defined.")
return pi
def choice_probs(cons, d, params_dict):
v = flow_util(cons, d, params_dict)
v_0 = flow_util(cons, 0, params_dict)
v_1 = flow_util(cons, 1, params_dict)
choice_prob = np.exp(v) / (np.exp(v_0) + np.exp(v_1))
return choice_prob
def m_util_aux(init_cond, params_dict, choice_0, nu, consumption_0):
"""Return the expected marginal utility for one realization of the wage shock."""
budget_0 = init_cond["assets_beginning_of_period"]
health_state_0 = init_cond["health"]
weighted_marginal = 0
for health_state_1 in [0, 1]:
for choice_1 in [0, 1]:
budget_1 = budget(
lagged_resources=budget_0,
lagged_consumption=consumption_0,
lagged_choice=choice_0,
wage=wage(nu=nu, params_dict=params_dict),
health=health_state_1,
params_dict=params_dict,
)
marginal_util = marginal_utility(budget_1, params_dict)
choice_prob = choice_probs(budget_1, choice_1, params_dict)
health_prob = prob_long_term_care_patient(
params_dict, health_state_0, health_state_1
)
weighted_marginal += choice_prob * health_prob * marginal_util
return weighted_marginal
def euler_rhs(init_cond, params_dict, draws, weights, choice_0, consumption):
discount_factor = params_dict["discount_factor"]
interest_factor = 1 + params_dict["interest_rate"]
rhs = 0
for index_draw, draw in enumerate(draws):
marg_util_draw = m_util_aux(init_cond, params_dict, choice_0, draw, consumption)
rhs += weights[index_draw] * marg_util_draw
return rhs * discount_factor * interest_factor
Comparison of DC-EGM algorithm and analytical solution¶
We now demonstrate the accuracy of the DC-EGM algorithm by substituting the calculated consumption policy into the Euler equation, showing that both sides are approximately equal.
For example, consider a specific state in the state space where initial health is \(D_0 = 0\) and initial wealth \(M_0\) is taken from the first (non-zero) entry on the wealth grid. Additionally, let \(d_0 = 0\).
choice_in_period_0 = 0
# Prepare dictionary for closed form solution
initial_condition = {"health": 0, "assets_beginning_of_period": 25}
state_dict = {
"ltc": initial_condition["health"],
"lagged_choice": 0,
"period": 0,
"assets_begin_of_period": initial_condition["assets_beginning_of_period"],
}
cons_calc, value = solved_model.policy_and_value_for_states_and_choices(
states=state_dict,
choices=choice_in_period_0,
)
Given the consumption policy calculated by dcegm, we now compare the left-hand side of the Euler equation (which is equal to the marginal utility of consumption) to its right-hand side.
# needed for computation of the integral
quad_points, quad_weights = roots_sh_legendre(5)
quad_draws = (
norm.ppf(quad_points) * params["income_shock_std"] + params["income_shock_mean"]
)
As expected, the consumption in the first period \(c_0\), as calculated by the DC-EGM algorithm, satisfies the Euler equation because both sides are (approximately) equal.
rhs = euler_rhs(
init_cond=initial_condition,
params_dict=params,
draws=quad_draws,
weights=quad_weights,
choice_0=choice_in_period_0,
consumption=cons_calc,
)
rhs
Array(0.08292194, dtype=float64)
lhs = marginal_utility(cons_calc, params)
lhs
Array(0.08292197, dtype=float64)