dcegm.interpolation.interpnd_regular¶

N-dimensional regular-grid interpolation for policy/value.

Assumptions: - Shared 1D wealth grid across all regular-grid combinations. - Child-state policy/value grids are flattened in regular dimensions:

(n_child_state_choices, n_continuous_combinations, n_wealth).

Child-state interpolation points are provided as continuous_state_child_states[name] with shape (n_child_state_choices, n_continuous_combinations).

Functions¶

`interpnd_policy_for_child_states_on_regular_grids`(...)	Interpolate policy, using value-based overwrite logic.
`interpnd_policy_and_value_for_child_states_on_regular_grids`(...)	Interpolate policy/value and apply consume-all overwrite.
`interpnd_value_for_child_states_on_regular_grids`(...)	Interpolate value and apply consume-all overwrite.
`_compute_consume_all_value`(→ jax.numpy.ndarray)
`_interp_policy_and_value_one_comb`(...)
`_interp_single_grid_one_comb`(→ jax.numpy.ndarray)
`_interp_wealth_for_corner`(→ jax.numpy.ndarray)
`_interp_regular_only_all`(→ jax.numpy.ndarray)	Interpolate values over regular dimensions only.
`_corner_linear_indices_and_weights`(...)
`_interp_regular_only`(→ jax.numpy.ndarray)
`_precompute_interp_objects`(→ Dict[str, jax.numpy.ndarray])	Precompute reusable interpolation objects for policy/value paths.
`_precompute_regular_indices_and_weights`(...)	Precompute low/high idx and weights for all regular child points.
`_take_last_axis`(arr, idx)	Take along last axis with batched indices.
`_compute_row_major_strides`(shape)
`_corner_table`(→ jax.numpy.ndarray)	Return binary corner table of shape `(2**n_dims, n_dims)`.

Module Contents¶

dcegm.interpolation.interpnd_regular.interpnd_policy_for_child_states_on_regular_grids(additional_continuous_state_grids: Dict[str, jax.numpy.ndarray], wealth_grid: jax.numpy.ndarray, policy_grid_child_states: jax.numpy.ndarray, value_grid_child_states: jax.numpy.ndarray, continuous_state_child_states: Dict[str, jax.numpy.ndarray], wealth_child_states: jax.numpy.ndarray, state_choice_child_states: Dict[str, Any], compute_utility: Callable, params: Dict[str, Any], discount_factor: float) → jax.numpy.ndarray¶

Interpolate policy, using value-based overwrite logic.

Returns shape (n_child_state_choices, n_continuous_combinations, n_wealth, n_quad_points).

dcegm.interpolation.interpnd_regular.interpnd_policy_and_value_for_child_states_on_regular_grids(additional_continuous_state_grids: Dict[str, jax.numpy.ndarray], wealth_grid: jax.numpy.ndarray, policy_grid_child_states: jax.numpy.ndarray, value_grid_child_states: jax.numpy.ndarray, continuous_state_child_states: Dict[str, jax.numpy.ndarray], wealth_child_states: jax.numpy.ndarray, state_choice_child_states: Dict[str, Any], compute_utility: Callable, params: Dict[str, Any], discount_factor: float) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]¶

Interpolate policy/value and apply consume-all overwrite.

If consume-all value dominates interpolated value at a point, overwrite policy with consume-all policy (=wealth point).

dcegm.interpolation.interpnd_regular.interpnd_value_for_child_states_on_regular_grids(additional_continuous_state_grids: Dict[str, jax.numpy.ndarray], wealth_grid: jax.numpy.ndarray, value_grid_child_states: jax.numpy.ndarray, continuous_state_child_states: Dict[str, jax.numpy.ndarray], wealth_child_states: jax.numpy.ndarray, state_choice_child_states: Dict[str, Any], compute_utility: Callable, params: Dict[str, Any], discount_factor: float) → jax.numpy.ndarray¶

Interpolate value and apply consume-all overwrite.

Returns shape (n_child_state_choices, n_continuous_combinations, n_wealth, n_quad_points).

dcegm.interpolation.interpnd_regular._compute_consume_all_value(expected_value_zero_savings: jax.numpy.ndarray, wealth_child_states: jax.numpy.ndarray, state_choice_child_states: Dict[str, Any], continuous_state_child_states: Dict[str, jax.numpy.ndarray], compute_utility: Callable, params: Dict[str, Any], discount_factor: float) → jax.numpy.ndarray¶

dcegm.interpolation.interpnd_regular._interp_policy_and_value_one_comb(policy_grid_one_child: jax.numpy.ndarray, value_grid_one_child: jax.numpy.ndarray, regular_low_idxs_one_comb: jax.numpy.ndarray, regular_high_idxs_one_comb: jax.numpy.ndarray, regular_low_weights_one_comb: jax.numpy.ndarray, regular_high_weights_one_comb: jax.numpy.ndarray, wealth_points_one_comb: jax.numpy.ndarray, wealth_low_idxs_one_comb: jax.numpy.ndarray, wealth_high_idxs_one_comb: jax.numpy.ndarray, strides: jax.numpy.ndarray, corner_table: jax.numpy.ndarray, wealth_grid: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]¶

dcegm.interpolation.interpnd_regular._interp_single_grid_one_comb(grid_one_child: jax.numpy.ndarray, corner_linear_idxs: jax.numpy.ndarray, corner_weights: jax.numpy.ndarray, wealth_points_one_comb: jax.numpy.ndarray, wealth_low_idxs_one_comb: jax.numpy.ndarray, wealth_high_idxs_one_comb: jax.numpy.ndarray, wealth_grid: jax.numpy.ndarray) → jax.numpy.ndarray¶

dcegm.interpolation.interpnd_regular._interp_wealth_for_corner(grid_row: jax.numpy.ndarray, wealth_points_one_comb: jax.numpy.ndarray, wealth_low_idxs_one_comb: jax.numpy.ndarray, wealth_high_idxs_one_comb: jax.numpy.ndarray, wealth_grid: jax.numpy.ndarray) → jax.numpy.ndarray¶

dcegm.interpolation.interpnd_regular._interp_regular_only_all(values_over_regular_grid_child_states: jax.numpy.ndarray, regular_low_idxs: jax.numpy.ndarray, regular_high_idxs: jax.numpy.ndarray, regular_low_weights: jax.numpy.ndarray, regular_high_weights: jax.numpy.ndarray, strides: jax.numpy.ndarray, corner_table: jax.numpy.ndarray) → jax.numpy.ndarray¶: Interpolate values over regular dimensions only.

dcegm.interpolation.interpnd_regular._corner_linear_indices_and_weights(regular_low_idxs_one_comb: jax.numpy.ndarray, regular_high_idxs_one_comb: jax.numpy.ndarray, regular_low_weights_one_comb: jax.numpy.ndarray, regular_high_weights_one_comb: jax.numpy.ndarray, strides: jax.numpy.ndarray, corner_table: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]¶

dcegm.interpolation.interpnd_regular._interp_regular_only(values_over_regular_grid: jax.numpy.ndarray, regular_low_idxs_one_comb: jax.numpy.ndarray, regular_high_idxs_one_comb: jax.numpy.ndarray, regular_low_weights_one_comb: jax.numpy.ndarray, regular_high_weights_one_comb: jax.numpy.ndarray, strides: jax.numpy.ndarray, corner_table: jax.numpy.ndarray) → jax.numpy.ndarray¶

dcegm.interpolation.interpnd_regular._precompute_interp_objects(additional_continuous_state_grids: Dict[str, jax.numpy.ndarray], continuous_state_child_states: Dict[str, jax.numpy.ndarray], wealth_grid: jax.numpy.ndarray, wealth_child_states: jax.numpy.ndarray) → Dict[str, jax.numpy.ndarray]¶: Precompute reusable interpolation objects for policy/value paths.

dcegm.interpolation.interpnd_regular._precompute_regular_indices_and_weights(additional_continuous_state_grids: Dict[str, jax.numpy.ndarray], continuous_state_child_states: Dict[str, jax.numpy.ndarray], state_names)¶

Precompute low/high idx and weights for all regular child points.

Returns arrays of shape (n_dims, n_child_state_choices, n_continuous_combinations).

dcegm.interpolation.interpnd_regular._take_last_axis(arr, idx)¶: Take along last axis with batched indices.

dcegm.interpolation.interpnd_regular._compute_row_major_strides(shape)¶

dcegm.interpolation.interpnd_regular._corner_table(n_dims: int) → jax.numpy.ndarray¶

Return binary corner table of shape (2**n_dims, n_dims).

Pseudo-code equivalent:

for corner in range(2**n_dims):

for dim in range(n_dims):: table[corner, dim] = (corner >> dim) & 1