Generalizing the mapping of indices with more than two state variable


I just have a question regarding the discrete dynamic programming package ddp. In general, there is a single index to represent the state space. Thus, if there are multiple state variables, one needs to map pairs of integers to one integer.

With two state variables, in the context of, say, the Aiyagari model, the mapping is as follows:

s_i = a_i *z_size + z_i

for indices on assets a_i and productivity z_i.

This can be inverted as follows:

a_i = s_i//z_size (integer division)
z_i = s_i% z_size (remainder)

How does this mapping generalize with more than two state variables?


I should also mention I am using Python, but that is somewhat beyond the point of the general logic of the question.

Hi @msilva913, this is a good question

In Julia the function you are after is called sub2ind.

I believe that the python equivalent is numpy.ravel_multi_index (docs here)

You would use it like this:

import numpy as np
s_i = np.ravel_multi_index((a_i, z_i), (a_size, s_size))

This will generalize to as many states as you wish.

As far as how this works, you basically generalize the algorithm you’ve stated above, thinking about it recursively. For example, suppose I have three states (a, z, k), three sizes (a_size, z_size, k_size), and three indices (a_i, z_i, k_i)

I can then do

s_i = (a_i * z_size + z_i) * k_size + k_i

which returns the same as np.ravel_multi_index((a_i, z_i, k_i), (a_size, z_size, k_size)):

In [16]: a_size, z_size, k_size = (3, 4, 5)

In [17]: a_i, z_i, k_i = (2, 3, 4)

In [19]: (a_i * z_size + z_i) * k_size + k_i
Out[19]: 59

In [21]: np.ravel_multi_index((a_i, z_i, k_i), (a_size, z_size, k_size))
Out[21]: 59

Thank you! The recursion is clear, and thanks for providing the command.