dcegm.interpolation.interp2d

Jax implementation of 2D interpolation.

Attributes

TRANSFORMATION_MAT

Functions

interp2d_policy_and_value_on_wealth_and_regular_grid(...)

Linear 2D interpolation on two grids where wealth has irregular spacing.

interp2d_value_on_wealth_and_regular_grid(...)

Interpolate the value function on a 2D grid.

interp2d_policy_on_wealth_and_regular_grid(...)

Interpolate the policy function on a 2D grid.

interp2d_policy(regular_points, wealth_points, ...)

Interpolate the policy function on a 2D grid.

interp2d_value_and_check_creditconstraint(...)

Interpolate the value function on a 2D grid and check for credit constraints.

find_grid_coords_for_interp(regular_grid, wealth_grid, ...)

Find the coordinates and indices of the 2D grid on which we interpolate.

interp2d(x_coords, y_coords, z_vals, x_new, y_new)

Perform linear 2D interpolation on an irregular quadrilateral.

determine_coordinates_in_unit_square(x, y, x_vec, y_vec)

Determine the relative coordinates of a point within the unit square.

compute_vertex_weights(x, y)

Compute interpolation weights for a point within the unit square.

Module Contents

dcegm.interpolation.interp2d.TRANSFORMATION_MAT
dcegm.interpolation.interp2d.interp2d_policy_and_value_on_wealth_and_regular_grid(regular_grid: jax.numpy.ndarray, wealth_grid: jax.numpy.ndarray, policy_grid: jax.numpy.ndarray, value_grid: jax.numpy.ndarray, regular_point_to_interp: jax.numpy.ndarray | float, wealth_point_to_interp: jax.numpy.ndarray | float, compute_utility: Callable, state_choice_vec: Dict[str, int], params: dict, discount_factor)

Linear 2D interpolation on two grids where wealth has irregular spacing.

The function interpolates by mapping the irregular quadrilateral defined by the grid coordinates into a canonical unit square to simplify the interpolation process.

Note, while the regular grid has a shape of (n_regular_grid_points,) and only contains the values of the regular grid, the irregular wealth grid is different for each regular grid point. The same is true for the policy and value grids. These three objects together define the 2D value and policy functions. Hence, each of them has shape (n_regular_grid_points, n_wealth_grid_points).

Parameters:

  • regular_grid (jnp.ndarray) – A 1d array of shape (n_regular_grid_points,) with the values of the regular grid.

  • wealth_grid (jnp.ndarray) – A 2d array of with the values of the irregular wealth grid over the regular grid points of shape (n_regular_grid_points, n_wealth_grid_points).

  • policy_grid (jnp.ndarray) – A 2d array with the policy values of shape (n_regular_grid_points, n_wealth_grid_points).

  • value_grid (jnp.ndarray) – A 2d array with the value function values of shape. (n_regular_grid_points, n_wealth_grid_points).

  • regular_point_to_interp (jnp.ndarray | float) – The regular point for which to interpolate the policy and value function.

  • wealth_point_to_interp (jnp.ndarray | float) – The wealth point for which to interpolate the policy and value function.

  • compute_utility (Callable) – User function to compute the utility of consumption.

  • params (dict) – A dictionary containing the model parameters.

Returns:

A tuple containing the interpolated values of the policy and

value function.

Return type:

tuple

dcegm.interpolation.interp2d.interp2d_value_on_wealth_and_regular_grid(regular_grid: jax.numpy.ndarray, wealth_grid: jax.numpy.ndarray, value_grid: jax.numpy.ndarray, regular_point_to_interp: jax.numpy.ndarray | float, wealth_point_to_interp: jax.numpy.ndarray | float, compute_utility: Callable, state_choice_vec: Dict[str, int], params: dict, discount_factor)

Interpolate the value function on a 2D grid.

This function interpolates the values of the value function at a specific point given by regular_point_to_interp and wealth_point_to_interp.

The function is useful in situtions where only the value but not the policy needs to be computed - such as maximum likelihood estimation.

Parameters:

  • regular_points (jnp.ndarray) – A 1d array of four elements representing the regular grid points used for interpolation.

  • wealth_points (jnp.ndarray) – A 1d array of four elements representing the wealth grid points used for interpolation.

  • policy_grid (jnp.ndarray) – A 2d array of policy function values with shape (n_regular_grid_points, n_wealth_grid_points).

  • coords_idxs (jnp.ndarray) – A 2d array of shape (2, 2) containing the indices of the (regular, wealth) grid where the interpolation is performed.

  • regular_point_to_interp (float | jnp.ndarray) – The regular grid point at which to interpolate.

  • wealth_point_to_interp (float | jnp.ndarray) – The wealth grid point at which to interpolate.

  • compute_utility (Callable) – User-defined function to compute the utility of consumption.

  • state_choice_vec (Dict[str, int]) – Dictionary specifying the state choices.

  • params (dict) – A dictionary containing model parameters.

Returns:

The interpolated value of the policy function at the given

(regular, wealth) point.

Return type:

float

dcegm.interpolation.interp2d.interp2d_policy_on_wealth_and_regular_grid(regular_grid: jax.numpy.ndarray, wealth_grid: jax.numpy.ndarray, policy_grid: jax.numpy.ndarray, regular_point_to_interp: jax.numpy.ndarray | float, wealth_point_to_interp: jax.numpy.ndarray | float)

Interpolate the policy function on a 2D grid.

This function interpolates the values of the policy function at a specific point given by regular_point_to_interp and wealth_point_to_interp.

The function is useful in situtions where only the policy but not the value needs to be computed - such as maximum likelihood estimation.

Parameters:

  • regular_points (jnp.ndarray) – A 1d array of four elements representing the regular grid points used for interpolation.

  • wealth_points (jnp.ndarray) – A 1d array of four elements representing the wealth grid points used for interpolation.

  • policy_grid (jnp.ndarray) – A 2d array of policy function values with shape (n_regular_grid_points, n_wealth_grid_points).

  • coords_idxs (jnp.ndarray) – A 2d array of shape (2, 2) containing the indices of the (regular, wealth) grid where the interpolation is performed.

  • regular_point_to_interp (float | jnp.ndarray) – The regular grid point at which to interpolate.

  • wealth_point_to_interp (float | jnp.ndarray) – The wealth grid point at which to interpolate.

Returns:

The interpolated value of the policy function at the given

(regular, wealth) point.

Return type:

jnp.ndarray | float

dcegm.interpolation.interp2d.interp2d_policy(regular_points, wealth_points, policy_grid, coords_idxs, regular_point_to_interp, wealth_point_to_interp)

Interpolate the policy function on a 2D grid.

This function interpolates the values of the policy function at a specific point given by regular_point_to_interp and wealth_point_to_interp.

Parameters:

  • regular_points (jnp.ndarray) – A 1d array of four elements representing the regular grid points used for interpolation.

  • wealth_points (jnp.ndarray) – A 1d array of four elements representing the wealth grid points used for interpolation.

  • policy_grid (jnp.ndarray) – A 2d array of policy function values with shape (n_regular_grid_points, n_wealth_grid_points).

  • coords_idxs (jnp.ndarray) – A 2d array of shape (2, 2) containing the indices of the (regular, wealth) grid where the interpolation is performed.

  • regular_point_to_interp (float | jnp.ndarray) – The regular grid point at which to interpolate.

  • wealth_point_to_interp (float | jnp.ndarray) – The wealth grid point at which to interpolate.

Returns:

The interpolated value of the policy function at the given

(regular, wealth) point.

Return type:

float

dcegm.interpolation.interp2d.interp2d_value_and_check_creditconstraint(regular_points, wealth_points, value_grid, coords_idxs, regular_point_to_interp, wealth_point_to_interp, compute_utility, wealth_min_unconstrained, value_at_zero_wealth, state_choice_vec, params, discount_factor)

Interpolate the value function on a 2D grid and check for credit constraints.

This function interpolates the value function at a specific point given by regular_point_to_interp and wealth_point_to_interp. It checks whether the point lies within a credit-constrained region and adjusts the interpolation accordingly using closed-form solutions for the value of consuming all wealth.

Parameters:

  • regular_points (jnp.ndarray) – A 1d array of four elements representing the regular grid points used for interpolation.

  • wealth_points (jnp.ndarray) – A 1d array of four elements representing the wealth grid points used for interpolation.

  • policy_grid (jnp.ndarray) – A 2d array of policy function values with shape (n_regular_grid_points, n_wealth_grid_points).

  • coords_idxs (jnp.ndarray) – A 2d array of shape (2, 2) containing the indices of the (regular, wealth) grid where the interpolation is performed.

  • regular_point_to_interp (float | jnp.ndarray) – The regular grid point at which to interpolate.

  • wealth_point_to_interp (float | jnp.ndarray) – The wealth grid point at which to interpolate.

  • compute_utility (callable) – A function to compute the utility of consumption.

  • wealth_min_unconstrained (jnp.ndarray) – A 1d array of minimum unconstrained wealth levels for each regular grid point.

  • value_at_zero_wealth (jnp.ndarray) – A 1d array of value function values at zero wealth for each regular grid point.

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

Returns:

The interpolated value of the value function at the given

(regular, wealth) point.

Return type:

float

dcegm.interpolation.interp2d.find_grid_coords_for_interp(regular_grid, wealth_grid, regular_point_to_interp, wealth_point_to_interp)

Find the coordinates and indices of the 2D grid on which we interpolate.

This function determines the four grid points and their indices for interpolation. It exploits the fact that one dimension of the 2D grid is regular. The wealth grid is irregular.

Parameters:

  • regular_grid (jnp.ndarray) – A 1d array of shape (n_regular_grid_points,) with the values of the regular grid.

  • wealth_grid (jnp.ndarray) – A 2d array with the values of the irregular wealth grid of shape (n_regular_grid_points, n_wealth_grid_points).

  • regular_point_to_interp (jnp.ndarray | float) – The regular point for which to interpolate the policy and value function.

  • wealth_point_to_interp (jnp.ndarray | float) – The wealth point for which to interpolate the policy and value function.

Returns:

  • regular_points (jnp.ndarray): A 1d array of shape (4,) with the four points

    of the regular grid.

  • wealth_points (jnp.ndarray): A 1d array of shape (4,) with the four points

    of the wealth grid.

  • coords_idxs (jnp.ndarray): A 2d array of shape (2, 2) containing the indices

    of the (regular, wealth) grid where the interpolation is performed.

The structure of the return arrays are the same in each row (or element):

  • The first element is the lower left point.

  • The second element is the lower right point.

  • The third element is the upper right point.

  • The fourth element is the upper left point.

Return type:

tuple

dcegm.interpolation.interp2d.interp2d(x_coords, y_coords, z_vals, x_new, y_new)

Perform linear 2D interpolation on an irregular quadrilateral.

This function maps the vertices of an irregular quadrilateral onto a canonical unit square and performs bilinear interpolation for a given point within this square.

Parameters:

  • x_coords (jnp.ndarray) – A 1d array of four elements representing the x-coordinates of the quadrilateral’s vertices.

  • y_coords (jnp.ndarray) – A 1d array of four elements representing the y-coordinates of the quadrilateral’s vertices.

  • z_vals (jnp.ndarray) – A 1d array of four elements representing the values at the vertices of the quadrilateral.

  • x_new (float) – The x-coordinate of the point where interpolation is required.

  • y_new (float) – The y-coordinate of the point where interpolation is required.

Returns:

The interpolated value at the point (x_new, y_new).

Return type:

float

dcegm.interpolation.interp2d.determine_coordinates_in_unit_square(x, y, x_vec, y_vec)

Determine the relative coordinates of a point within the unit square.

This function computes the relative coordinates of a point (x, y) within the unit square, using transformed basis vectors x_vec and y_vec. It maps the point of an irregular quadrilateral point into a normalized space to ease interpolation.

Parameters:

  • x (float) – The x-coordinate of the point to be mapped.

  • y (float) – The y-coordinate of the point to be mapped.

  • x_vec (jnp.ndarray) – A 1d array of four elements representing the transformed x-coordinates of the quadrilateral’s vertices.

  • y_vec (jnp.ndarray) – A 1d array of two elements representing the transformed y-coordinates of the quadrilateral’s vertices.

Returns:

A tuple (x_rel, y_rel) representing the relative coordinates of the

point within the unit square. x_rel is the relative position along the x-axis, and y_rel is the relative position along the y-axis.

Return type:

tuple

dcegm.interpolation.interp2d.compute_vertex_weights(x, y)

Compute interpolation weights for a point within the unit square.

This function calculates the weights for bilinear interpolation based on the relative position of a point (x, y) within a canonical unit square. These weights are used to determine the contribution of each vertex of the square to the interpolated value.

Parameters:

  • x (float) – The relative x-coordinate of the point within the unit square.

  • y (float) – The relative y-coordinate of the point within the unit square.

Returns:

A 1d array of four weights corresponding to the contributions of

the lower left, lower right, upper right, and upper left vertices of the unit square, respectively.

Return type:

jnp.ndarray