Enable JAX persistent caching via EconPizzaConfig class

.github/workflows/continuous-integration.yml

@@ -21,12 +21,12 @@ jobs:
       fail-fast: False
       matrix:
         os: [ubuntu-latest, macos-latest, windows-latest]
-        python-version: ['3.9', '3.10', '3.11']
+        python-version: ['3.10', '3.11', '3.12']
 
     steps:
-    - uses: actions/checkout@v3
+    - uses: actions/checkout@master
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v3
+      uses: actions/setup-python@master
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies

README.rst

@@ -1,18 +1,19 @@
-econpizza:
+econpizza
 =========
-
-**Solve nonlinear heterogeneous agent models using automatic differentiation**
-
-.. image:: https://img.shields.io/badge/GitHub-gboehl%2Feconpizza-blue.svg?style=flat
+.. |badge0| image:: https://img.shields.io/badge/GitHub-gboehl%2Feconpizza-blue.svg?style=flat
     :target: https://github.com/gboehl/econpizza
-.. image:: https://github.com/gboehl/econpizza/actions/workflows/continuous-integration.yml/badge.svg
+.. |badge1| image:: https://github.com/gboehl/econpizza/actions/workflows/continuous-integration.yml/badge.svg
     :target: https://github.com/gboehl/econpizza/actions
-.. image:: https://readthedocs.org/projects/econpizza/badge/?version=latest
+.. |badge2| image:: https://readthedocs.org/projects/econpizza/badge/?version=latest
     :target: http://econpizza.readthedocs.io/en/latest/?badge=latest
-.. image:: https://badge.fury.io/py/econpizza.svg
+.. |badge3| image:: https://badge.fury.io/py/econpizza.svg
     :target: https://badge.fury.io/py/econpizza
 
-**Econpizza** is a framework to solve and simulate *fully nonlinear* perfect foresight models, with or without heterogeneous agents.
+|badge0| |badge1| |badge2| |badge3|
+
+**Solve nonlinear heterogeneous agent models using automatic differentiation**
+
+Econpizza is a framework to solve and simulate **fully nonlinear** perfect foresight models, with or without heterogeneous agents.
 The package implements the solution method proposed in `HANK on Speed: Robust Nonlinear Solutions using Automatic Differentiation <https://gregorboehl.com/live/hank_speed_boehl.pdf>`_ *(Gregor Boehl, 2023, SSRN No. 4433585)*.
 It allows to specify and solve nonlinear macroeconomic models quickly in a simple, high-level fashion and provides generic and robust routines for steady state search.
 
@@ -22,11 +23,15 @@ The method extends the `Sequence-Space Jacobian <https://github.com/shade-econ/s
 To solve models with representative agents a shooting methods similar to Laffargue (1990), Boucekkine (1995) and Juillard (1996) is implemented. It is faster and more reliable than the extended path method in dynare due to the use of automatic differentiation for the efficient jacobian decompositions during each Newton-step. Nonlinear perfect-foresight transition dynamics can - even for large-scale nonlinear models with several occassionally binding constraints - be computed in less than a second.
 
 The package builds heavily on `automatic differentiation <https://en.wikipedia.org/wiki/Automatic_differentiation>`_ via `JAX <https://jax.readthedocs.io/en/latest/notebooks/quickstart.html>`_.
+There is a `presentation <https://gregorboehl.com/revealjs/adh/index.html>`_ on how this works behind the szenes.
 
 
 Documentation
 -------------
 
+▶ `User guide <https://econpizza.readthedocs.io/en/stable/index.html>`_
+▶ `Quickstart tutorial <https://econpizza.readthedocs.io/en/stable/tutorial/quickstart.html>`_
+
 Installing the `repository version <https://pypi.org/project/econpizza/>`_ from PyPi is as simple as typing
 
 .. code-block:: bash
@@ -35,11 +40,6 @@ Installing the `repository version <https://pypi.org/project/econpizza/>`_ from
 
 in your terminal or Anaconda Prompt.
 
-Guides and tutorials are provided on ReadTheDocs:
-
- * `User guide <https://econpizza.readthedocs.io/en/stable/index.html>`_
- * `Quickstart tutorial <https://econpizza.readthedocs.io/en/stable/tutorial/quickstart.html>`_
-
 Citation
 --------
 .. code-block:: bibtex

docs/guide/solution.rst

@@ -21,9 +21,9 @@ Any additional argument will be passed on to the specific Newton method. For mod
 
 .. autofunction:: econpizza.utilities.newton.newton_for_jvp
 
-For models with representative agents, the Newton method is :meth:`econpizza.utilities.newton.newton_for_banded_jac`:
+For models with representative agents, the Newton method is :meth:`econpizza.utilities.newton.newton_for_tridiag_jac`:
 
-.. autofunction:: econpizza.utilities.newton.newton_for_banded_jac
+.. autofunction:: econpizza.utilities.newton.newton_for_tridiag_jac
 
 If ``use_solid_solver`` is set to `True`, the Newton method `newton_jax_jit <https://grgrjax.readthedocs.io/en/latest/#grgrjax.newton_jax_jit>`_ from the `grgrjax <https://grgrjax.readthedocs.io>`_ package is used.
 

docs/guide/steady_state.rst

@@ -7,11 +7,13 @@ Upon failure, the function tries to be as informative as possible. If the search
 
 .. note::
 
-   A classic complaint is "**The Jacobian contains NaNs**". This is usually due to numerical errors somewhere along the way. While the package tries to provide more information about where the error occurred, a good idea is to follow JAX's hints on `how to debug NaNs <https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#debugging-nans>`_.
+   A classic complaint is "**Function returns NaNs**" or "**The Jacobian contains NaNs**". This is usually either due infinite/non-existent derivatives or invalid values (e.g. `sqrt(-1)` or `1/0`) along the way. While the package tries to provide more information about where the error occurred, a good idea is to follow JAX's hints on `how to debug NaNs <https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#debugging-nans>`_.
 
 .. tip::
 
    * A common gotcha for heterogeneous agent models is that the distribution contains negative values. The routine will be informative about that. This is usually due to rather excessive interpolation outside the grid and can often be fixed by using a grid with larger minimum/maximum values.
+   * To avoid NaNs, ensure that parameter values produce valid results on an unbounded domain. E.g., instead of declaring a parameter `nu` when calling `sqrt(nu)`, declare a parameter `nu_log` and define `nu = exp(nu_log)` in the `aux_equations` block. This ensures that `nu` is always positive.
+   * A known problem is that for steady state values set to zero (either as fixed value or initial guess), the Jacobian matrix may contain NaNs. This is due to how JAX internally calculates Jacobians via Jacobian-vector-products. It then often helps to set these values to `1e-8` instead of `0`.
    * The steady state procedure is based on the `pseudoinverse <https://en.wikipedia.org/wiki/Moore%E2%80%93Penrose_inverse>`_ of the Jacobian. If the procedure fails, it will try to tell you the rank of the Jacobian and the number of degrees of freedom. More degrees of freedom than the Jacobian rank normally implies that you should fix more steady state values and vice versa.
    * If the desired precision is not reached, the error message will tell you in which equation the maximum error did arise. You can use the ``equations`` key to get the list of equations (as strings), e.g. ``print(model['equations'][17])`` to get the equation with index 17.
 

docs/guide/the_yaml.rst

@@ -23,7 +23,8 @@ The YAML files follow a simple structure:
 4. optionally provide auxiliary equations that are not directly part of the nonlinear system
 5. optionally provide initial guesses for all other steady state values and parameters
 
-I will first briefly discuss the YAML of the small scale *representative* agents NK model `from the quickstart tutorial <../tutorial/quickstart.ipynb>`_ and then turn to a more complex HANK model.
+I will first briefly discuss `the YAML <https://github.com/gboehl/econpizza/blob/master/econpizza/examples/nk.yml>`_ of the small scale *representative* agents NK model `from the quickstart tutorial <../tutorial/quickstart.ipynb>`_ and then turn to a more complex HANK model.
+A collection of examples is provided `with the package <https://github.com/gboehl/econpizza/tree/master/econpizza/examples>`_.
 
 YAML: representative agent models
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -32,14 +33,14 @@ The GitHub version of the YAML file for the small scale NK model can be found `h
 
 .. code-block:: yaml
 
-    variables: [y, c, pi, r, rn, beta, w, chi]
+    variables: [y, c, pi, r, rn, beta, w]
     shocks: [e_beta]
 
 Note that it is not necessary to define shocks. You can also simply set the initial values of any (exogenous) state.
 
 .. code-block:: yaml
 
-    parameters: [ theta, psi, phi_pi, phi_y, rho, h, eta, rho_beta, chi ]
+    parameters: [ theta, psi, phi_pi, phi_y, rho, h, sigma_l, rho_beta, chi ]
 
 Use the ``parameters`` block to define any *parameters*. Parameters are treated the same as variables, but they are time invariant. During steady state search they are treated exactly equally. For this reason their values are provided in the ``steady_state`` block.
 
@@ -53,7 +54,7 @@ The second block (``definitions``) defines general definitions and imports, whic
 .. code-block:: yaml
 
     equations:
-        ~ w = chi*(c - h*cLag)*y**eta  # labor supply
+        ~ w = chi*(c - h*cLag)*y**sigma_l  # labor supply
         ~ 1 = r*betaPrime*(c - h*cLag)/(cPrime - h*c)/piPrime  # euler equation
         ~ psi*(pi/piSS - 1)*pi/piSS = (1-theta) + theta*w + psi*betaPrime*(c-h*cLag)/(cPrime-h*c)*(piPrime/piSS - 1)*piPrime/piSS*yPrime/y  # Phillips curve
         ~ c = (1-psi*(pi/piSS - 1)**2/2)*y  # market clearing
@@ -70,23 +71,23 @@ Note that you need one equation for each variable defined in ``variables``.
         fixed_values:
             # parameters
             theta: 6.  # demand elasticity
-            psi: 96  # price adjustment costs
-            phi_pi: 4  # monetary policy rule coefficient #1
-            phi_y: 1.5  # monetary policy rule coefficient #2
+            psi: 60  # price adjustment costs
+            phi_pi: 1.5  # monetary policy rule coefficient #1
+            phi_y: 0.1  # monetary policy rule coefficient #2
             rho: .8  # interest rate smoothing
-            h: .44  # habit formation
-            eta: .33  # inverse Frisch elasticity
+            h: .74  # habit formation
+            sigma_l: 2  # inverse Frisch elasticity
             rho_beta: .9  # autocorrelation of discount factor shock
 
             # steady state values
-            beta: 0.9984
+            beta: 0.995
             y: .33
             pi: 1.02^.25
 
-        init_guesses: # the default initial guess is always 1.1
+        init_guesses: # the default initial guess for values not specified here is always 0.95
             chi: 6
 
-Finally, the ``steady_state`` block allows to fix parameters and, if desired, some steady state values, and provide initial guesses for others. Note that the default initial guess for any variable/parameter not specified here will be ``1.1``.
+Finally, the ``steady_state`` block allows to fix parameters and, if desired, some steady state values, and provide initial guesses for others. Note that the default initial guess for any variable/parameter not specified here will be ``0.95``.
 
 
 YAML: heterogeneous agent models

docs/tutorial/boehl_hommes.rst

@@ -43,7 +43,7 @@ The following example shows how to use the shooting method on the simple New Key
     x, _, flag = mod.find_path_shooting(state.values())
 
     # plotting
-    for i,v in enumerate(mod['variables']):
+    for i,v in enumerate(mod.var_names):
 
         plt.figure()
         plt.plot(x[:,i])
@@ -63,14 +63,14 @@ Lets go for a second, numerically more challenging example: the chaotic rational
     _ = mod.solve_stst()
 
     # choose an interesting initial state
-    state = np.zeros(len(mod['variables']))
+    state = np.zeros(len(mod.var_names))
     state[:-1] = [.1, .2, 0.]
 
     # solve and simulate
     x, _, flag = ep.find_path_shooting(mod, state, T=500, max_horizon=1000, tol=1e-5)
 
     # plotting
-    for i,v in enumerate(mod['variables']):
+    for i,v in enumerate(mod.var_names):
 
         plt.figure()
         plt.plot(x[:,i])