Under the hood

The functional representations of the economic model are written dynamically during parsing/loarding (in econpizza/parser/__init__.py).

import econpizza as ep
example_hank = ep.examples.hank

mod = ep.load(example_hank)

WARNING:jax._src.xla_bridge:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

(load:) Parsing done.

The model instance is a dictionary, containing all the informations of the model. For instance, it contains the dynamically created functions as strings:

mod['func_strings'].keys()

dict_keys(['func_backw', 'func_eqns'])

The function func_backw corresponds to function \(W(\cdot)\) from the paper and func_eqns is \(f(\cdot)\). The other functions are static.

Lets inspect \(f\):

print(mod['func_strings']['func_eqns'])

def func_eqns(XLag, X, XPrime, XSS, shocks, pars, distributions=[], decisions_outputs=[]):


 (BLag, betaLag, CLag, divLag, nLag, piLag, RLag, RnLag, RrLag, RstarLag, taxLag, Top10ALag, Top10CLag, wLag, yLag, y_prodLag, zLag, ) = XLag

 (B, beta, C, div, n, pi, R, Rn, Rr, Rstar, tax, Top10A, Top10C, w, y, y_prod, z, ) = X

 (BPrime, betaPrime, CPrime, divPrime, nPrime, piPrime, RPrime, RnPrime, RrPrime, RstarPrime, taxPrime, Top10APrime, Top10CPrime, wPrime, yPrime, y_prodPrime, zPrime, ) = XPrime

 (BSS, betaSS, CSS, divSS, nSS, piSS, RSS, RnSS, RrSS, RstarSS, taxSS, Top10ASS, Top10CSS, wSS, ySS, y_prodSS, zSS, ) = XSS

 (sigma_c, sigma_l, theta, psi, phi_pi, phi_y, rho, rho_beta, rho_r, rho_z, ) = pars

 (e_beta, e_rstar, e_z, ) = shocks

 (dist, ) = distributions

 (a, c, ) = decisions_outputs

 # NOTE: summing over the first two dimensions e and a, but not the time dimension (dimension 2)
 # `dist` here corresponds to the dist *at the beginning of the period*
 aggr_a = jnp.sum(dist*a, axis=(0,1))
 aggr_c = jnp.sum(dist*c, axis=(0,1))
 # calculate consumption and wealth share of top-10%
 top10c = 1 - percentile(c, dist, .9)
 top10a = 1 - percentile(a, dist, .9)

 root_container0 = C  - ( aggr_c)
 root_container1 = Top10C  - ( top10c)
 root_container2 = Top10A  - ( top10a)
 root_container3 = n  - ( y_prod/z)
 root_container4 = div  - ( -w*n + (1 - psi*(pi/piSS - 1)**2/2)*y_prod)
 root_container5 = y  - ( (1 - psi*(pi/piSS - 1)**2/2)*y_prod)
 root_container6 = psi*(pi/piSS - 1)*pi/piSS  - ( (1-theta) + theta*w + psi*piPrime/R*(piPrime/piSS - 1)*piPrime/piSS*y_prodPrime/y_prod)
 root_container7 = tax  - ( (Rr-1)*BLag)
 root_container8 = Rr  - ( RLag/pi)
 root_container9 = Rn  - ( (Rstar*((pi/piSS)**phi_pi)*((y/yLag)**phi_y))**(1-rho)*RnLag**rho)
 root_container10 = R  - ( maximum(1, Rn))
 root_container11 = C  - ( y)
 root_container12 = B  - ( aggr_a)
 root_container13 = n**sigma_l  - ( w)
 root_container14 = beta  - ( betaSS*(betaLag/betaSS)**rho_beta*exp(e_beta))
 root_container15 = Rstar  - ( RstarSS*(RstarLag/RstarSS)**rho_r*exp(e_rstar))
 root_container16 = z  - ( zSS*(zLag/zSS)**rho_z*exp(e_z))

 return jnp.array([root_container0, root_container1, root_container2, root_container3, root_container4, root_container5, root_container6, root_container7, root_container8, root_container9, root_container10, root_container11, root_container12, root_container13, root_container14, root_container15, root_container16]).T.ravel()

This function is then automatically compiled and the callable can be found in model['context']:

mod['context']['func_eqns']

<function econpizza.parser.func_eqns(XLag, X, XPrime, XSS, shocks, pars, distributions=[], decisions_outputs=[])>

The model['context'] itself contans the name space in which all model functions and definitions are evaluated. This may be useful for debugging:

mod['context'].keys()

dict_keys(['__name__', '__doc__', '__package__', '__loader__', '__spec__', '__path__', '__file__', '__cached__', '__builtins__', 'yaml', 're', 'os', 'sys', 'tempfile', 'jax', 'jaxlib', 'jnp', 'iu', 'deepcopy', 'copy', 'getmembers', 'isfunction', 'jax_print', 'het_agent_base_funcs', 'build_functions', 'write_dynamic_functions', 'func_forw_generic', 'func_forw_stst_generic', 'compile_func_basics_str', 'compile_backw_func_str', 'get_forw_funcs', 'compile_eqn_func_str', 'checks', 'func_pre_stst', 'check_if_defined', 'check_dublicates', 'check_determinancy', 'check_initial_values', 'check_shapes', 'check_if_compiled', 'grids', 'dists', 'interp', 'cached_mdicts', 'cached_models', 'd2jnp', '_load_as_module', 'parse', '_eval_strs', '_parse_external_functions_file', '_initialize_context', '_initialize_cache', '_load_external_functions_file', '_compile_init_values', '_define_subdict_if_absent', '_define_function', '_get_pre_stst_mapping', 'compile_stst_inputs', 'load', 'log', 'exp', 'sqrt', 'max', 'min', 'egm_init', 'egm_step', 'interpolate', 'transfers', 'maximum', 'percentile', 'skills_grid', 'skills_stationary', 'skills_transition', 'a_grid', 'func_backw', 'func_forw', 'func_forw_stst', 'func_eqns', 'sigma_c', 'sigma_l', 'theta', 'psi', 'phi_pi', 'phi_y', 'rho', 'rho_beta', 'rho_r', 'rho_z', 'y', 'y_prod', 'C', 'pi', 'beta', 'B', 'w', 'n', 'div', 'z', 'Rstar', 'Rr', 'Rn', 'R', 'tax', 'WaPrime', 'init_run'])