search

jessegrabowski / gEconpy

A collection of tools for working with DSGE models in python, inspired by the R package gEcon
https://geconpy.readthedocs.io/en/latest/index.html
GNU General Public License v3.0
21 stars 3 forks source link

njit the estimation pipeline end-to-end #4

Open jessegrabowski opened 1 year ago

jessegrabowski commented 1 year ago

The main gEconModel is currently not jittable and cannot be rewritten as a jitclass because of its heavy dependence on sympy.

During estimation, however, I want the main loop of steady state -> log linearization -> BK check -> Kalman Filter to be entirely jit-compiled. This is currently accomplished using a set of functions in estimation.estimation_utilities that extract the necessary components from the model and jit them all separately.

I think it would be better to instead "compile" the model to a jitclass that holds all the necessary functionality. This would also clean up the codebase, which is currently too fragmented into cluttered "utility" files.

jessegrabowski commented 1 year ago

This was partially implemented by #29 accidentally. The new function numba_tools.utilities.numba_lambdify makes very fast numba code from sympy equations. In addition, they can be pickled. That means steady state and log linearization are fully jitted (but only for solver="cyclic_reduction", gensys still needs to be jitted). BK check is jitted in estimation.estimation_utilities.check_bk_condition, but not in the model itself -- this is code duplication and should be unified. The kalman filters are also almost fully jitted. The last barrier to fast sampling is the build_and_solve function, but I'd like to rebuild the whole sampling pipeline if possible.