fix formatting

gboehl · gboehl · commit f0c56c225dba · 2025-02-26T12:57:37.000+01:00
diff --git a/docs/guide/installation.rst b/docs/guide/installation.rst
@@ -11,5 +11,5 @@ Installing the `repository version <https://pypi.org/project/econpizza/>`_ from
 
 in your terminal or Anaconda Prompt.
 
-The current version is supporting Python versions 3.9 to 3.11 on Windows, Mac and Linux.
+The current version is supporting the latest three minor Python versions on Windows, Mac and Linux.
 The `changelog and release history <https://github.com/gboehl/econpizza/releases>`_ can be found on GitHub.
diff --git a/docs/tutorial/boehl_hommes.rst b/docs/tutorial/boehl_hommes.rst
@@ -2,9 +2,9 @@
 Boehl-Hommes method
 -------------------
 
-The package also contains an alternative "shooting" method much aligned to the one introduced in Boehl & Hommes (2021). In the original paper we use this method to solve for chaotic asset price dynamics. The method can be understood as a policy function iteration where the initial state is the only fixed grid point and all other grid points are chosen endogenously (as in a "reverse" EGM) to map the expected trajectory.
+The package also contains an alternative "shooting" method much aligned to the one introduced in Boehl & Hommes (2021). In the original paper we use this method to solve for chaotic asset price dynamics. The method can be understood as an extension to Fair & Taylor (1983) and is similar to a  policy function iteration where the initial state is the only fixed grid point and all other grid points are chosen endogenously (as in a "reverse" EGM) to map the expected trajectory.
 
-The main advantage (in terms of robustness) over the stacking method comes from exploiting the property that most determined perfect foresight models are a contraction mapping both, forward and backwards. The model is given by
+Assume the model is given by
 
 .. code-block::
 
@@ -17,39 +17,9 @@ We iterate on the expected trajectory itself instead of the policy function. We
    d f(x_{t-1}, x_t, x_{t+1} ) < d x_{t-1},
    d f(x_{t-1}, x_t, x_{t+1} ) < d x_{t+1}.
 
-This is also the weakness of the method: not every DSGE model (that is determined in the Blanchard-Kahn sense) is such backward-and-forward contraction. In most cases the algorithm converges anyways, but convergence is not guaranteed.
+This is also the weakness of the method: not every DSGE model (that is determined in the Blanchard-Kahn sense) is such backward-and-forward contraction. Thus, unless you are interested in chaotic dynamics, the standard "stacking" method is to be prefered.
 
-The following example shows how to use the shooting method on the simple New Keynesian model.
-
-.. code-block:: python
-
-    import numpy as np
-    import matplotlib.pyplot as plt
-    import econpizza as ep
-    from econpizza import example_nk
-
-    # load the example.
-    # example_nk is nothing else but the path to the yaml, hence you could also use `filename = 'path_to/model.yaml'`
-    mod = ep.load(example_nk)
-    # solve for the steady state
-    _ = mod.solve_stst()
-
-    # get the steady state as an initial state
-    state = mod['stst'].copy()
-    # increase the discount factor by one percent
-    state['beta'] *= 1.02
-
-    # simulate the model
-    x, _, flag = mod.find_path_shooting(state.values())
-
-    # plotting
-    for i,v in enumerate(mod.var_names):
-
-        plt.figure()
-        plt.plot(x[:,i])
-        plt.title(v)
-
-Lets go for a second, numerically more challenging example: the chaotic rational expectations model of Boehl & Hommes (2021)
+The following example shows how to use the shooting method to a numerically challenging example: the chaotic rational expectations model of Boehl & Hommes (2021). The YAML for this model can be found `here <https://github.com/gboehl/econpizza/blob/master/econpizza/examples/bh.yml>`_.
 
 .. code-block:: python
 
@@ -81,3 +51,8 @@ This will give you boom-bust cycles in asset pricing dynamics:
 .. image:: https://github.com/gboehl/econpizza/blob/master/docs/p_and_n.png?raw=true
   :width: 800
   :alt: Dynamics of prices and fractions
+
+Below, find the documentation for the additional function:
+
+.. autofunction:: econpizza.PizzaModel.find_path_shooting
+
diff --git a/docs/tutorial/hank1.ipynb b/docs/tutorial/hank1.ipynb
diff --git a/docs/tutorial/hank2.ipynb b/docs/tutorial/hank2.ipynb
diff --git a/econpizza/__version__.py b/econpizza/__version__.py
@@ -1,3 +1,3 @@
 # -*- coding: utf-8 -*-
 
-__version__ = '0.6.5'
+__version__ = '0.6.6'
diff --git a/econpizza/examples/dsge.yml b/econpizza/examples/dsge.yml
@@ -76,7 +76,7 @@ steady_state:
         rho: 0.8 # persistence in (notional) nominal interest rate
         omega_p: 0.44 # coefficient on steady state inflation in price indexation
         omega_w: 0.66 # coefficient on steady state wage inflation in wage indexation
-        iota_w: 1 # degree of dowards nominal wage rigidity
+        iota_w: 0 # degree of dowards nominal wage rigidity
         elb: 1 # position of ELB
         rho_beta: 0.9 # persistence of discount factor shock
         rho_z: 0.9 # persistence of technology shocks
diff --git a/econpizza/parser/__init__.py b/econpizza/parser/__init__.py
@@ -348,6 +348,7 @@ def load(
     if model.get('decisions'):
         decisions_outputs = model['decisions']['outputs']
         decisions_inputs = model['decisions']['inputs']
+        check_for_lags(model['decisions']['calls'], evars)
         model['func_strings']["func_backw"] = compile_backw_func_str(
             evars, par_names, shocks, decisions_inputs, decisions_outputs, model['decisions']['calls'])
         _define_function(model['func_strings']
diff --git a/econpizza/parser/build_generic_functions.py b/econpizza/parser/build_generic_functions.py
@@ -170,7 +170,7 @@ def get_stacked_func_het_agents(func_backw, func_forw, func_eqns, stst, wfSS, ho
 
     # build partials of output functions
     backwards_sweep_local = jax.tree_util.Partial(
-        backwards_sweep, stst=stst, wfSS=wfSS, horizon=horizon, func_backw=partial_backw)
+        backwards_sweep, stst=stst, wfSS=wfSS, func_backw=partial_backw)
     forwards_sweep_local = jax.tree_util.Partial(
         forwards_sweep, horizon=horizon, func_forw=partial_forw)
     final_step_local = jax.tree_util.Partial(
diff --git a/econpizza/parser/checks.py b/econpizza/parser/checks.py
@@ -129,3 +129,10 @@ def check_if_compiled(model, horizon, pars, stst):
         return jnp.allclose(model['cache']['pars'], pars)
     except:
         return False
+
+
+def check_for_lags(calls, evars):
+    for v in evars:
+        if v+'Lag' in calls:
+            raise Exception(
+                f"`{v}Lag` in decisions calls detected. For efficiency reasons, the use of lagged values is not supported here. This can be circumvented by defining an auxilliary variable, e.g. `{v}_lagged = {v}Lag`.")
diff --git a/econpizza/parser/het_agent_base_funcs.py b/econpizza/parser/het_agent_base_funcs.py
@@ -34,8 +34,10 @@ def _backwards_step(carry, i):
     return (wf, X, shocks, func_backw, stst, pars), (wf, decisions_output)
 
 
-def backwards_sweep(x: Array, x0: Array, shocks: Array, pars: Array, stst: Array, wfSS: Array, horizon: int, func_backw: Callable, return_wf=False) -> Array:
+def backwards_sweep(x: Array, x0: Array, shocks: Array, pars: Array, stst: Array, wfSS: Array, func_backw: Callable, return_wf=False) -> Array:
 
+    # get horizon from input size
+    horizon = len(x)//len(stst) + 1
     X = jnp.hstack((x0, x, stst)).reshape(horizon+1, -1).T
 
     _, (wf_storage, decisions_output_storage) = jax.lax.scan(
diff --git a/econpizza/solvers/shooting.py b/econpizza/solvers/shooting.py
@@ -42,26 +42,24 @@ def find_path_shooting(
     shock : tuple, optional
         shock in period 0 as in `(shock_name_as_str, shock_size)`
     horizon : int, optional
-        number of periods to simulate
+        number of periods to actually simulate. This is NOT the expectations horizon of agents. Defaults to 30
     init_path : array, optional
-        a first guess on the trajectory. Normally not necessary
+        an initial guess on the trajectory. This should not be necessary
     max_horizon : int, optional
-        number of periods until the system is assumed to be back in the steady state. A good idea to set this corresponding to the respective problem. Note that a horizon too far away may cause the accumulation of numerical errors.
+        number of periods until the system is assumed to be back in the steady state. A good idea to set this corresponding to the respective problem. Note that a horizon too far away may cause the accumulation of numerical errors. Defaults to 200
     max_loops : int, optional
-        number of repetitions to iterate over the whole trajectory. Should eventually be high.
-    max_iterations : int, optional
-        number of iterations. Default is `max_horizon`. It should not be lower than that (and will raise an error). Normally it should not be higher, better use `max_loops` instead.
+        number of repetitions to iterate over the whole trajectory. Should eventually be high. Defaults to 100
+    max_iter : int, optional
+        number of iterations. Defaults to `max_horizon`. It should not be lower than that (and will raise an error). Normally it should not be higher, better use `max_loops` instead.
     tol : float, optional
-        convergence criterion
+        convergence criterion, defaults to `1e-5`
     verbose : bool, optional
-        degree of verbosity. 0/`False` is silent
+        degree of verbosity. 0/`False` is silent. Defualts to `True`
 
     Returns
     -------
     x_fin : array
         array of the trajectory
-    x_lin : array or None
-        array of the trajectory based on the linear model. Will return None if the linear model is unknown
     fin_flag : int
         error code
     """
diff --git a/econpizza/solvers/solve_linear.py b/econpizza/solvers/solve_linear.py
@@ -49,6 +49,8 @@ def find_path_linear(self, shock=None, init_state=None, pars=None, horizon=200,
     nvars = len(self["var_names"])
     pars = d2jnp((pars if pars is not None else self["pars"]))
     shocks = self.get("shocks") or ()
+    distSS = jnp.array(self['steady_state'].get('distributions'))
+    doSS = [jnp.array(d)[..., None] for d in list(self['steady_state']['decisions'].values())]
     x_stst = jnp.ones((horizon + 1, nvars)) * stst
 
     # deal with shocks
@@ -63,17 +65,19 @@ def find_path_linear(self, shock=None, init_state=None, pars=None, horizon=200,
         derivatives = get_stst_derivatives(
             self, nvars, pars, stst, x_stst, zero_shocks, horizon, verbose)
 
-        # accumulate steady stat jacobian
+        # accumulate steady state jacobian
         get_stst_jacobian(self, derivatives, horizon, nvars, verbose)
+        jac_f2xLag = derivatives[0][0]
+        self['cache']['jacobian_f2x0'] = jnp.zeros(((horizon-1)*nvars, nvars)).at[:nvars].set(jac_f2xLag[...,0])
         write_cache(self, horizon, pars, stst)
 
-    jacobian = self['cache']['jac']
+    # get jacobians
+    jacobian = self['cache']['jac_factorized']
+    jacobian_f2x0 = self['cache']['jacobian_f2x0']
+    combined_sweep = self['context']['combined_sweep']
 
-    x0 -= stst
-    x = - \
-        jax.scipy.linalg.solve(
-            jacobian[nvars:, nvars:], jacobian[nvars:, :nvars] @ x0)
-    x = jnp.hstack((x0, x)).reshape(-1, nvars) + stst
+    x = -jax._src.lax.linalg.lu_solve(*jacobian[0], jacobian_f2x0 @ (x0-stst), 0)[jacobian[1]]
+    x = jnp.vstack((x0,x.reshape(-1, nvars) + stst, stst))
 
     if verbose:
         duration = time.time() - st
diff --git a/econpizza/solvers/stacking.py b/econpizza/solvers/stacking.py
@@ -44,15 +44,15 @@ def find_path_stacking(
     init_dist : array, optional
         initial distribution, defaults to the steady state distribution
     init_guess : array, optional
-        initial guess on the sequence trajectory, defaults to the steady state
+        initial guess on the sequence trajectory (the solution), defaults to the steady state
     pars : dict, optional
-        alternative parameters. Warning: do only change those parameters that are invariant to the steady state.
+        alternative parameters. Warning: it will only be effective to change those parameters that are invariant to the steady state.
     horizon : int, optional
-        number of periods until the system is assumed to be back in the steady state. Defaults to ``200``
+        number of periods until the system is assumed to be back in the steady state. Finding trajectories with a low horizon is quicker, but chosing a larger horizon may be necessary if the simulated time series do not converge. Defaults to ``200``
     use_solid_solver : bool, optional
-        calculate the full jacobian and use a standard Newton method. Defaults to ``False``
+        calculate the full jacobian instead of the iterative approximation, and use a standard Newton method. This usually takes a very long time. Defaults to ``False``
     skip_jacobian : bool, optional
-        whether to skip the calculation of the steady state sequence space Jacobian. If True, the last cached Jacobian will be used. Defaults to ``False``
+        whether to skip the calculation of the steady state sequence space Jacobian. If True, the last cached Jacobian will be used. Only use this if you know what you are doing. Defaults to ``False``
     verbose : bool, optional
         degree of verbosity. ``0``/``False`` is silent. Defaults to ``False``
     raise_errors : bool, optional
@@ -89,7 +89,8 @@ def find_path_stacking(
         'distributions')
     dist0 = jnp.array(init_dist if init_dist is not None else jnp.nan)
     x_stst = jnp.ones((horizon + 1, nvars)) * stst
-    x_init = init_guess if init_guess is not None else x_stst.at[0].set(x0)
+    x_init = init_guess if init_guess is not None else x_stst
+    x_init = x_init.at[0].set(x0)
 
     # deal with shocks if any
     shock_series = jnp.zeros((horizon-1, len(shocks)))
diff --git a/econpizza/tools.py b/econpizza/tools.py
@@ -44,6 +44,11 @@ def traverse_dict(d, path, value=None):
         return eval(str(dd))
 
 
+def print_formated_found_values(model):
+    _ = [print(f"        {k}: {v}") for k,v in model['steady_state']['found_values'].items()]
+    return
+
+
 def anneal_stst(mdict, dict_path, target_value, max_sequence=10, **kwargs):
     """Anneal steady state by iteratively updating initial guesses. Assumes that the current model dictionary provides valid results.
 
@@ -96,7 +101,7 @@ def anneal_stst(mdict, dict_path, target_value, max_sequence=10, **kwargs):
 
     # print final values
     print('(anneal_stst:) Success! Values are:\n')
-    [print(f"        {k}: {v}") for k,v in current_model['steady_state']['found_values'].items()]
+    print_formated_found_values(current_model)
     return current_model, mdict
 
 
diff --git a/econpizza/utilities/newton.py b/econpizza/utilities/newton.py
@@ -35,8 +35,9 @@ def inner_iteration_cond(carry):
 def jvp_while_body(carry):
     (x, y, _, _, cnt, fev, _), statics = carry
     (jvp_func, jacobian, maxit, relaxation, nsteps, tol, factor, verbose) = statics
-    # first iteration
+    # apply newton
     x -= relaxation*y
+    # first iteration
     f, _ = jvp_func(x, jnp.zeros_like(x))
     f = lu_solve(*jacobian[0], f, 0)[jacobian[1]]
     # other iterations
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,6 @@
 # ubuntu seems to have issues with jax v0.4.6. Recheck when v0.4.7 is out
-jax>=0.4.13
-jaxlib>=0.4.13
+jax>=0.4.13,<0.5.1
+jaxlib>=0.4.13,<0.5.1
 grgrjax>=0.4.6
 scipy
 pyyaml
diff --git a/setup.py b/setup.py
@@ -34,8 +34,8 @@
         'linear': ['grgrlib>=0.1.22'],
     },
     install_requires=[
-        "jax>=0.4.13",
-        "jaxlib>=0.4.13",
+        "jax>=0.4.13,<0.5.1",
+        "jaxlib>=0.4.13,<0.5.1",
         "grgrjax>=0.4.6",
         "pyyaml",
         "scipy",

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,3 @@`
`1`	`1`	`# -- coding: utf-8 --`
`2`	`2`
`3`		`-__version__ = '0.6.5'`
	`3`	`+__version__ = '0.6.6'`