dcegm.egm.interpolate_marginal_utility

Functions

interpolate_value_and_marg_util(...)	Interpolate value and policy for all child states and compute marginal utility.
interp1d_value_and_marg_util_for_state_choice(...)	Interpolate value and policy for given child state and compute marginal utility.
_interpolate_value_and_marg_util_2d_irregular(...)
_interpolate_value_and_marg_util_nd_regular(...)
_get_continuous_states_next(→ Dict[str, jax.numpy.ndarray])
_compute_nd_marginal_utility(→ jax.numpy.ndarray)	Compute marginal utility pointwise on ND interpolation output via vmaps.
interp2d_value_and_marg_util_for_state_choice(...)	Interpolate value and policy for given child state and compute marginal utility.

Module Contents

dcegm.egm.interpolate_marginal_utility.interpolate_value_and_marg_util(model_funcs, state_choice_vec: Dict[str, int], continuous_grids_info: Dict[str, Any], cont_grids_next_period: Dict[str, jax.numpy.ndarray], endog_grid_child_state_choice: jax.numpy.ndarray, policy_child_state_choice: jax.numpy.ndarray, value_child_state_choice: jax.numpy.ndarray, child_state_idxs: jax.numpy.ndarray, continuous_state_space, params: Dict[str, float], upper_envelope_method: str) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

Interpolate value and policy for all child states and compute marginal utility.

Parameters:

• compute_marginal_utility (callable) – User-defined function to compute the agent’s marginal utility of consumption.

• compute_utility (callable) – Function for calculating the utility of consumption.

• state_choice_vec (dict) – Dictionary containing the state and choice of the agent.

• assets_beginning_of_next_period (jnp.ndarray) – 2d array of shape (n_quad_stochastic, n_grid_wealth,) containing the agent’s beginning of period wealth.

• endog_grid_child_state_choice (jnp.ndarray) – 1d array containing the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• policy_child_state_choice (jnp.ndarray) – 1d array containing the corresponding policy function values of the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• value_child_state_choice (jnp.ndarray) – 1d array containing the corresponding value function values of the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• has_second_continuous_state (bool) – Boolean indicating whether the model features a second continuous state variable. If False, the only continuous state variable is consumption/savings.

• params (dict) – Dictionary containing the model parameters.

Returns:

• value_interp (jnp.ndarray): 2d array of shape (n_wealth_grid, n_income_shocks)

  containing the interpolated value function.

• marg_util_interp (jnp.ndarray): 2d array of shape (n_wealth_grid, n_income_shocks)

  containing the interpolated marginal utilities for each wealth level and income shock.

Return type:

tuple

dcegm.egm.interpolate_marginal_utility.interp1d_value_and_marg_util_for_state_choice(compute_marginal_utility: Callable, compute_utility: Callable, state_choice_vec: Dict[str, int], assets_beginning_of_next_period: jax.numpy.ndarray, endog_grid_child_state_choice: jax.numpy.ndarray, policy_child_state_choice: jax.numpy.ndarray, value_child_state_choice: jax.numpy.ndarray, params: Dict[str, float], discount_factor: float, use_dj_interpolation: bool) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

Interpolate value and policy for given child state and compute marginal utility.

Parameters:

• compute_marginal_utility (callable) – User-defined function to compute the agent’s marginal utility of consumption.

• compute_utility (callable) – Function for calculating the utility of consumption.

• state_choice_vec (dict) – Dictionary containing the state and choice of the agent.

• assets_beginning_of_next_period (jnp.ndarray) – 2d array of shape (n_quad_stochastic, n_grid_wealth,) containing the agent’s beginning of period wealth.

• endog_grid_child_state_choice (jnp.ndarray) – 1d array containing the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• policy_child_state_choice (jnp.ndarray) – 1d array containing the corresponding policy function values of the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• value_child_state_choice (jnp.ndarray) – 1d array containing the corresponding value function values of the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• has_second_continuous_state (bool) – Boolean indicating whether the model features a second continuous state variable. If False, the only continuous state variable is consumption/savings.

• params (dict) – Dictionary containing the model parameters.

Returns:

• marg_utils (jnp.ndarray): 2d array of shape (n_wealth_grid, n_income_shocks)

  containing the interpolated marginal utilities for each wealth level and income shock.

• value_interp (jnp.ndarray): 2d array of shape (n_wealth_grid, n_income_shocks)

  containing the interpolated value function.

Return type:

tuple

dcegm.egm.interpolate_marginal_utility._interpolate_value_and_marg_util_2d_irregular(compute_marginal_utility: Callable, compute_utility: Callable, state_choice_vec: Dict[str, int], continuous_grids_info: Dict[str, Any], cont_grids_next_period: Dict[str, jax.numpy.ndarray], child_state_idxs: jax.numpy.ndarray, wealth_child_states: jax.numpy.ndarray, endog_grid_child_state_choice: jax.numpy.ndarray, policy_child_state_choice: jax.numpy.ndarray, value_child_state_choice: jax.numpy.ndarray, params: Dict[str, float], discount_factor: float) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

dcegm.egm.interpolate_marginal_utility._interpolate_value_and_marg_util_nd_regular(compute_marginal_utility: Callable, compute_utility: Callable, state_choice_vec: Dict[str, int], continuous_grids_info: Dict[str, Any], cont_grids_next_period: Dict[str, jax.numpy.ndarray], child_state_idxs: jax.numpy.ndarray, wealth_child_states: jax.numpy.ndarray, endog_grid_child_state_choice: jax.numpy.ndarray, policy_child_state_choice: jax.numpy.ndarray, value_child_state_choice: jax.numpy.ndarray, params: Dict[str, float], discount_factor: float) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

dcegm.egm.interpolate_marginal_utility._get_continuous_states_next(cont_grids_next_period: Dict[str, jax.numpy.ndarray]) → Dict[str, jax.numpy.ndarray]

dcegm.egm.interpolate_marginal_utility._compute_nd_marginal_utility(compute_marginal_utility: Callable, policy_interp: jax.numpy.ndarray, state_choice_child_states: Dict[str, Any], continuous_state_child_states: Dict[str, jax.numpy.ndarray], params: Dict[str, float]) → jax.numpy.ndarray

Compute marginal utility pointwise on ND interpolation output via vmaps.

dcegm.egm.interpolate_marginal_utility.interp2d_value_and_marg_util_for_state_choice(compute_marginal_utility: Callable, compute_utility: Callable, state_choice_vec: Dict[str, int], continuous_state_space: Dict[str, jax.numpy.ndarray], assets_beginning_of_next_period: jax.numpy.ndarray, continuous_state_beginning_of_next_period: jax.numpy.ndarray, endog_grid_child_state_choice: jax.numpy.ndarray, policy_child_state_choice: jax.numpy.ndarray, value_child_state_choice: jax.numpy.ndarray, params: Dict[str, float], discount_factor: float) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray]

Interpolate value and policy for given child state and compute marginal utility.

Parameters:

• compute_marginal_utility (callable) – User-defined function to compute the agent’s marginal utility of consumption.

• compute_utility (callable) – Function for calculating the utility of consumption.

• state_choice_vec (dict) – Dictionary containing the state and choice of the agent.

• assets_beginning_of_next_period (jnp.ndarray) – 2d array of shape (n_quad_stochastic, n_grid_wealth,) containing the agent’s beginning of period wealth.

• endog_grid_child_state_choice (jnp.ndarray) – 1d array containing the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• policy_child_state_choice (jnp.ndarray) – 1d array containing the corresponding policy function values of the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• value_child_state_choice (jnp.ndarray) – 1d array containing the corresponding value function values of the endogenous wealth grid of the child state/choice pair. Shape (n_grid_wealth,).

• has_second_continuous_state (bool) – Boolean indicating whether the model features a second continuous state variable. If False, the only continuous state variable is consumption/savings.

• params (dict) – Dictionary containing the model parameters.

Returns:

• marg_utils (jnp.ndarray): 2d array of shape (n_wealth_grid, n_income_shocks)

  containing the interpolated marginal utilities for each wealth level and income shock.

• value_interp (jnp.ndarray): 2d array of shape (n_wealth_grid, n_income_shocks)

  containing the interpolated value function.

Return type:

tuple