Time Iteration¶

Time iteration is the core nonlinear policy-iteration algorithm in Dolo for solving dynamic stochastic general equilibrium (DSGE) models, optionally subject to occasionally binding inequality constraints.

1. Mathematical Formulation¶

Classical Structural Form¶

In its classical structural formulation, a discrete-time continuous-state economic model is represented by:

\[\begin{aligned} s_t & = g\left(m_{t-1}, s_{t-1}, x_{t-1}, m_t \right) \\ 0 & = \mathbb{E}_t \left[ f\left(m_t, s_t, x_t, m_{t+1}, s_{t+1}, x_{t+1} \right) \right] \quad \perp \quad \underline{x}(m_t, s_t) \le x_t \le \overline{x}(m_t, s_t) \end{aligned}\]

where: - $m_t$ is the vector of exogenous shocks, - $s_t$ is the vector of endogenous states, - $x_t$ is the vector of controls (decisions), - $g$ is the state transition law, - $f$ is the arbitrage equation (Euler residuals), - $\underline{x}, \overline{x}$ are lower and upper bounds on controls, with the complementarity condition $\perp$.

The Unified $\tau$ Transition Operator¶

Following the design of Dolo.jl, dolo.py generalizes this formulation by abstracting the state transition and stochastic process into a unified transition kernel $\tau$.

Instead of treating exogenous Markov chains, VAR(1) processes, and i.i.d. quadrature shocks through disparate control flows, the next state distribution is written universally as:

\[S_{t+1} \sim \tau(S_t, x_t)\]

where a state $S$ is represented as a QualifiedPoint (or QP) bundling: 1. loc (Locator / Address): The discrete regime index or computational address in the discretized domain (e.g. regime index $i_m$ for Markov chains, or coordinate indices for table lookups). 2. val (Evaluated Payload): The continuous economic state coordinates $[m, s]$ passed into equations.

Given a policy rule $x = \varphi(S)$, the conditional transition $\tau(S_t, x_t)$ yields a set of discrete branches with weights and target qualified states:

\[\tau(S_t, x_t) = \left\{ \big(w_k, \, S^{(k)}_{t+1}\big) \right\}_{k=1}^K, \quad \sum_{k=1}^K w_k = 1\]

The arbitrage equation is formulated universally as:

\[0 = \mathbb{E}_{\tau} \left[ f\left(S_t, x_t, S_{t+1}, \varphi(S_{t+1})\right) \right] = \sum_{(w_k, S'_k) \in \tau(S_t, x_t)} w_k \cdot f\left(S_t, x_t, S'_k, \varphi(S'_k)\right)\]

This abstraction ensures that the Euler residuals and solution algorithms are identical whether the model features pure continuous states, discrete Markov regimes, or mixed exogenous processes.

2. Discrete Model & Core Abstractions¶

`DiscretizedModel`¶

A DiscretizedModel bundles the continuous model specification, the computational grid $\mathcal{G}$, and the discrete transition operator $\tau$:

\[\langle \text{Model}, \; \mathcal{G}, \; \tau \rangle\]

from dolo import yaml_import
from dolo.numeric import DiscretizedModel

model = yaml_import("examples/models/rbc.yaml")
dmodel = DiscretizedModel(model)

It provides: - dmodel.initial_guess(): Returns a GArray initialized from steady-state calibration. - dmodel.lower_bounds(), dmodel.upper_bounds(): Returns GArray bounds. - dmodel.arbitrage(...): Evaluates arbitrage equations across regimes and batches. - dmodel.transition_matrix(dr): Computes the full row-stochastic Markov transition matrix $P$ on grid nodes. - dmodel.ergodic_distribution(dr): Computes the stationary distribution $\mu P = \mu$.

`GArray` (Grid-Bound Fields)¶

A GArray represents a discrete field of values living on a grid $\mathcal{G}$. It exhibits a powerful duality: - Geometric Indexing: Slice by discrete regime or subgrid, e.g., controls[i_m]. - Vectorized Solver Interface: Directly provides ravel() (flat 1D view for SciPy/Newton solvers) and unravel(flat_x) to reconstruct structured fields without manual tensor reshaping. - Linear Algebra: Implements vectorized addition, scaling, and $L_\infty$ distance norm (ga.norm()).

3. The Time Iteration Algorithm¶

The algorithm proceeds as follows:

Initial Guess: Start with an initial policy field $x^{(0)} = \varphi^{(0)} \in \mathcal{G}$ (typically the calibrated steady-state controls).
Policy Update (Euler Equation Solve): Given future policy $\varphi^{(n)}$, solve for current period controls $\varphi^{(n+1)}(S)$ at each grid node $S \in \mathcal{G}$: $$0 = \sum_{(w_k, S'_k) \in \tau(S, x)} w_k \cdot f\left(S, x, S'_k, \varphi^{(n)}(S'_k)\right) \quad \perp \quad \underline{x}(S) \le x \le \overline{x}(S)$$ This step uses a vectorized serial Newton-Raphson or Non-Linear Complementarity Problem (NCP) solver.
Policy Acceleration (Optional Multi-Step Lookahead): To accelerate convergence, an intermediate policy evaluation step can be applied using Neumann series or GMRES linear solves (similar to Howard improvement in value iteration): $$\varphi^{(n+1)} \leftarrow \varphi^{(n+1)} + (I - J_x)^{-1} \cdot R\big(\varphi^{(n+1)}\big)$$
Convergence Check: Compute the residual norm $\epsilon_{n+1} = \|R(\varphi^{(n+1)})\|_\infty$ and the policy step error $\eta_{n+1} = \|\varphi^{(n+1)} - \varphi^{(n)}\|_\infty$:
If $\epsilon_{n+1} < \tau_\epsilon$ and $\eta_{n+1} < \tau_\eta$, terminate and return the converged decision rule $\varphi^*$.
Otherwise, update $\varphi^{(n)} \leftarrow \varphi^{(n+1)}$ and return to Step 2.

4. Usage & Options¶

Basic Usage¶

from dolo import yaml_import, time_iteration

# Load model
model = yaml_import("examples/models/rbc.yaml")

# Solve using time iteration
sol = time_iteration(model)

# Access decision rule
dr = sol.dr
print("Converged:", sol.converged)
print("Iterations:", sol.iterations)

Accelerated Time Iteration¶

For models with high persistence or slow policy adjustment, acceleration drastically reduces the number of outer iterations:

# Enable Neumann policy acceleration
sol = time_iteration(model, accelerate=True)

# Custom acceleration configuration
sol = time_iteration(
    model,
    accelerate={
        "linsolve": "neumann",  # or "gmres"
        "maxit": 50,
        "wait": 5,             # wait 5 iterations before starting acceleration
        "tol": 1e-8,
    }
)

5. API Reference¶

Finds a global solution for model using backward time-iteration.

Parameters¶

model : Model, ModelV2, or DiscretizedModel model to be solved engine / backend : str ("auto", "jax", "gpu", "cpu", "vectorized") computational engine for residual and Jacobian evaluations dr0 : DecisionRule, optional initial guess for the decision rule tol_eps : float (1e-8) tolerance on Euler equation residuals ||F(x)|| tol_eta : float (1e-6) tolerance on policy update step ||x_{k+1} - x_k|| maxit : int (500) maximum number of outer time-iteration loops inner : dict, optional options for the inner pointwise Euler solver: dict(maxit=10, max_backsteps=10, tol=1e-8) accelerate : bool or dict (False) if True or a dict, accelerates policy iterations using a preconditioned Neumann or GMRES step on (I - L) dx = (x_1 - x_0) after a warm-up phase. e.g. dict(wait=10, linsolve="neumann", maxit=50, safeguarded=False) interpolation : str ("cubic") interpolation method ("cubic" or "linear") verbose : bool (True) display iteration progress trace : bool (False) record intermediate decision rules

Source code in src/dolo/algos/time_iteration.py

def time_iteration(
    model: Model,
    *,
    engine: str = "auto",
    backend: str = None,
    dr0: DecisionRule = None,
    tol_eps: float = 1e-8,
    tol_eta: float = 1e-6,
    maxit: int = 500,
    inner: dict = None,
    accelerate=False,
    interpolation: str = "cubic",
    verbose: bool = True,
    details: bool = True,
    trace: bool = False,
    ignore_constraints: bool = False,
    hook=None,
    # Backward compatibility:
    dprocess=None,
    **kwargs,
) -> TimeIterationResult:
    """Finds a global solution for ``model`` using backward time-iteration.

    Parameters
    ----------
    model : Model, ModelV2, or DiscretizedModel
        model to be solved
    engine / backend : str ("auto", "jax", "gpu", "cpu", "vectorized")
        computational engine for residual and Jacobian evaluations
    dr0 : DecisionRule, optional
        initial guess for the decision rule
    tol_eps : float (1e-8)
        tolerance on Euler equation residuals ||F(x)||
    tol_eta : float (1e-6)
        tolerance on policy update step ||x_{k+1} - x_k||
    maxit : int (500)
        maximum number of outer time-iteration loops
    inner : dict, optional
        options for the inner pointwise Euler solver:
        dict(maxit=10, max_backsteps=10, tol=1e-8)
    accelerate : bool or dict (False)
        if True or a dict, accelerates policy iterations using a preconditioned
        Neumann or GMRES step on (I - L) dx = (x_1 - x_0) after a warm-up phase.
        e.g. dict(wait=10, linsolve="neumann", maxit=50, safeguarded=False)
    interpolation : str ("cubic")
        interpolation method ("cubic" or "linear")
    verbose : bool (True)
        display iteration progress
    trace : bool (False)
        record intermediate decision rules
    """
    # Absorb legacy aliases and flat arguments via kwargs
    if "ftol" in kwargs:
        tol_eps = kwargs.pop("ftol")
    if "xtol" in kwargs:
        tol_eta = kwargs.pop("xtol")
    if "tol" in kwargs:
        tol_eta = kwargs.pop("tol")
    if "interp_method" in kwargs:
        interpolation = kwargs.pop("interp_method")
    if "improve" in kwargs:
        accelerate = kwargs.pop("improve")

    with_complementarities = kwargs.pop("with_complementarities", None)
    if with_complementarities is not None:
        pass
    else:
        with_complementarities = not ignore_constraints

    inner_maxit = kwargs.pop("inner_maxit", None)
    max_backsteps = kwargs.pop("max_backsteps", None)

    # Inner solver configuration
    inner_cfg = {"maxit": 10, "max_backsteps": 10, "tol": 1e-6}
    if inner is not None:
        inner_cfg.update(inner)
    if inner_maxit is not None:
        inner_cfg["maxit"] = inner_maxit
    if max_backsteps is not None:
        inner_cfg["max_backsteps"] = max_backsteps

    # Acceleration configuration
    acc_flag = accelerate
    if acc_flag is False or acc_flag is None:
        acc_cfg = None
    elif acc_flag is True:
        acc_cfg = {
            "wait": 10,
            "linsolve": "neumann",
            "maxit": 50,
            "tol": 1e-8,
            "safeguarded": False,
        }
    elif isinstance(acc_flag, dict):
        acc_cfg = {
            "wait": 10,
            "linsolve": "neumann",
            "maxit": 50,
            "tol": 1e-8,
            "safeguarded": False,
            **acc_flag,
        }
        if acc_cfg.get("linsolve") == "iti":
            acc_cfg["linsolve"] = "neumann"
    else:
        acc_cfg = None

    if trace:
        trace_details = []
    else:
        trace_details = None

    # Construct or adopt DiscretizedModel
    if isinstance(model, DiscretizedModel):
        dmodel = model
    else:
        grid_options = kwargs.pop("grid_options", None)
        dprocess_options = kwargs.pop("dprocess_options", {})
        if dprocess is not None:
            dprocess_options["dprocess"] = dprocess
        dmodel = DiscretizedModel(
            model, grid_options=grid_options, dprocess_options=dprocess_options
        )

    grid = dmodel.grid
    dprocess = dmodel.dprocess

    n_ms = dmodel.n_regimes
    n_im = dprocess.n_inodes(0)

    x0 = dmodel.model.calibration["controls"]
    parms = dmodel.parms
    n_x = len(x0)
    n_s = len(dmodel.model.symbols["states"])

    endo_grid = grid["endo"]
    exo_grid = grid["exo"]

    mdr = DecisionRule(
        exo_grid, endo_grid, dprocess=dprocess, interp_method=interpolation
    )
    if acc_cfg is not None:
        mdr_filt = DecisionRule(
            exo_grid, endo_grid, dprocess=dprocess, interp_method=interpolation
        )
        jres_alloc = numpy.zeros((n_ms, n_im, endo_grid.nodes.shape[0], n_x, n_x))
        S_ij_alloc = numpy.zeros((n_ms, n_im, endo_grid.nodes.shape[0], n_s))
    else:
        mdr_filt = None

    s = mdr.endo_grid.nodes
    N = s.shape[0]

    controls_ga = dmodel.initial_guess()
    if dr0 is None:
        pass
    else:
        if isinstance(dr0, AlgoResult):
            dr0 = dr0.dr
        try:
            for i_m in range(n_ms):
                controls_ga[i_m] = dr0(i_m, s)
        except Exception:
            for i_m in range(n_ms):
                m = dprocess.node(i_m)
                controls_ga[i_m] = dr0(m, s)

    f = dmodel.f
    g = dmodel.g

    if with_complementarities:
        lb_ga = dmodel.lower_bounds()
        ub_ga = dmodel.upper_bounds()
        if lb_ga is None or ub_ga is None:
            with_complementarities = False
            lb_flat = None
            ub_flat = None
            lb = None
            ub = None
        else:
            lb_flat = lb_ga.ravel().reshape((-1, n_x))
            ub_flat = ub_ga.ravel().reshape((-1, n_x))
            lb = lb_ga.data
            ub = ub_ga.data
    else:
        with_complementarities = False
        lb = None
        ub = None
        lb_flat = None
        ub_flat = None

    sh_c = controls_ga.shape
    controls_0 = controls_ga.ravel().reshape((-1, n_x))

    # Engine setup
    eng = _resolve_engine(dmodel, engine, backend)
    use_jax = eng in ("jax", "gpu", "cpu") and dmodel.n_regimes == 1
    if use_jax:
        import jax
        import jax.numpy as jnp

        target_dev = None
        if eng == "gpu":
            gpus = jax.devices("gpu")
            if not gpus:
                raise RuntimeError("GPU requested (engine='gpu'), but no CUDA device was found.")
            target_dev = gpus[0]
        elif eng == "cpu":
            target_dev = jax.devices("cpu")[0]

        cr = _get_or_compile_jax_residual(dmodel, interp_method=interpolation, diff_mode="fwd")
        s_nodes_jax = jnp.asarray(dmodel.endo_nodes, dtype=jnp.float64)
        if target_dev is not None:
            s_nodes_jax = jax.device_put(s_nodes_jax, target_dev)
    else:
        cr = None
        target_dev = None

    itprint = IterationsPrinter(
        ("N", int),
        ("Error", float),
        ("Gain", float),
        ("Time", float),
        ("nit", int),
        verbose=verbose,
    )
    itprint.print_header("Start Time Iterations.")

    t1 = time.time()
    err_0 = numpy.nan
    err_step = 10.0
    err_res = 10.0
    it = 0
    verbit = verbose == "full"

    while (err_step > tol_eta or err_res > tol_eps) and it < maxit:
        it += 1
        t_start = time.time()

        mdr.set_values(controls_ga)

        if trace:
            trace_details.append({"dr": copy.deepcopy(mdr)})

        if use_jax:
            coeffs_jax = cr.prefilter(controls_ga)
            if target_dev is not None:
                coeffs_jax = jax.device_put(coeffs_jax, target_dev)

            def dfn_jax(x_val):
                x_jnp = jnp.asarray(x_val, dtype=jnp.float64)
                if target_dev is not None:
                    x_jnp = jax.device_put(x_jnp, target_dev)
                r_jnp, dr_jnp = cr.F_and_dF_vmap(s_nodes_jax, x_jnp, coeffs_jax)
                if target_dev is not None:
                    r_jnp.block_until_ready()
                    dr_jnp.block_until_ready()
                return [numpy.asarray(r_jnp), numpy.asarray(dr_jnp)]

            fn = lambda x_val: dfn_jax(x_val)[0]
            dfn = dfn_jax
            res_0, _ = dfn(controls_0)
            err_res = abs(res_0).max()
        else:
            fn = (
                lambda x_val: euler_residuals_tau(dmodel, controls_ga.unravel(x_val), mdr)
                .ravel()
                .reshape((-1, n_x))
            )
            dfn = SerialDifferentiableFunction(fn)

            res_0 = fn(controls_0)
            err_res = abs(res_0).max()

        if hook:
            hook()

        # Inner Euler solver
        if with_complementarities:
            [controls, nit] = ncpsolve(
                dfn,
                lb_flat,
                ub_flat,
                controls_0,
                verbose=verbit,
                maxit=inner_cfg["maxit"],
            )
        else:
            [controls, nit] = serial_newton(
                dfn, controls_0, verbose=verbit, maxit=inner_cfg["maxit"]
            )

        new_controls_ga = controls_ga.unravel(controls)
        err_step = (new_controls_ga - controls_ga).norm()
        controls_ga = new_controls_ga
        err_SA = err_step / err_0 if not numpy.isnan(err_0) and err_0 > 0 else numpy.nan
        err_0 = err_step

        # Policy acceleration step
        if acc_cfg is not None and it > acc_cfg["wait"]:
            d = (controls - controls_0).reshape(sh_c)
            acc_method = acc_cfg.get("linsolve", "neumann")
            acc_maxit = acc_cfg.get("maxit", 50)
            acc_tol = acc_cfg.get("tol", 1e-8)

            if use_jax and not with_complementarities:
                from dolo.numeric.residuals import dF_2

                cand_ga = controls_ga.unravel(controls)
                mdr.set_values(cand_ga)
                _, dres_c = dfn(controls)
                J_c = dres_c[0] if dres_c.ndim == 4 else dres_c

                L_op = dF_2(
                    dmodel,
                    cand_ga,
                    dr=mdr,
                    engine=eng,
                    interp_method=interpolation,
                    device=target_dev,
                )
                L_prec = (-1.0) * L_op.preconditioned(J_c)
                tot = L_prec.solve(d, linsolve=acc_method, tol=acc_tol, maxit=acc_maxit)
                cand_controls_0 = controls_0 + numpy.asarray(tot).reshape((-1, n_x))
            else:
                # Evaluate derivatives at candidate controls
                mdr.set_values(controls.reshape(sh_c))
                res_c, dres_c = dfn(controls)
                res_c = res_c.reshape(sh_c)
                dres_c = dres_c.reshape((sh_c[0], N, n_x, n_x))

                jres_alloc[...] = 0.0
                S_ij_alloc[...] = 0.0
                junk, jres, fut_S = euler_residuals_acc(
                    f,
                    g,
                    s,
                    controls.reshape(sh_c),
                    mdr,
                    dprocess,
                    parms,
                    diff=False,
                    with_jres=True,
                    set_dr=False,
                    jres=jres_alloc,
                    S_ij=S_ij_alloc,
                )

                if with_complementarities:
                    res_c, dres_c, jres = smooth(
                        res_c, dres_c, jres, controls.reshape(sh_c) - lb
                    )
                    res_c[...] *= -1
                    dres_c[...] *= -1
                    jres[...] *= -1
                    res_c, dres_c, jres = smooth(
                        res_c, dres_c, jres, ub - controls.reshape(sh_c), pos=-1.0
                    )
                    res_c[...] *= -1
                    dres_c[...] *= -1
                    jres[...] *= -1

                jres[...] *= -1.0

                for i_m in range(n_ms):
                    for j_m in range(n_im):
                        M_mat = jres[i_m, j_m, :, :, :]
                        X_mat = dres_c[i_m, :, :, :].copy()
                        solve_tensor(X_mat, M_mat)

                if acc_method == "gmres":
                    L = Operator(jres, fut_S, mdr_filt)
                    L.addid = True
                    try:
                        sol_gmres, info = scipy.sparse.linalg.gmres(
                            L, d.ravel(), rtol=acc_tol, atol=0.0, maxiter=acc_maxit
                        )
                    except TypeError:
                        sol_gmres, info = scipy.sparse.linalg.gmres(
                            L, d.ravel(), tol=acc_tol, maxiter=acc_maxit
                        )
                    tot = sol_gmres.reshape(d.shape)
                else:
                    tot, *rest = neumann_solve(
                        d,
                        jres,
                        fut_S,
                        mdr_filt,
                        tol=acc_tol,
                        maxit=acc_maxit,
                        verbose=(verbose == "full"),
                    )

                cand_controls_0 = controls_0 + tot.reshape((-1, n_x))

            if acc_cfg.get("safeguarded", False):
                cand_res = fn(cand_controls_0)
                if abs(cand_res).max() < err_res:
                    controls_0 = cand_controls_0
                else:
                    controls_0 = controls
            else:
                controls_0 = cand_controls_0
        else:
            controls_0 = controls

        controls_ga = controls_ga.unravel(controls_0)

        t_finish = time.time()
        elapsed = t_finish - t_start

        itprint.print_iteration(
            N=it, Error=err_step, Gain=err_SA, Time=elapsed, nit=nit
        )

        if err_step < tol_eta and it > 1:
            break

    mdr.set_values(controls_ga)

    if trace:
        trace_details.append({"dr": copy.deepcopy(mdr)})

    itprint.print_finished()

    if not details:
        return mdr

    converged = bool(err_step < tol_eta or err_res < tol_eps)

    return TimeIterationResult(
        mdr,
        it,
        with_complementarities,
        dprocess,
        converged,
        tol_eta,
        err_step,
        None,
        trace_details,
        err_res,
    )

Bases: AlgoResult

Source code in src/dolo/algos/results.py

@dataclass
class TimeIterationResult(AlgoResult):
    dr: object
    iterations: int
    complementarities: bool
    dprocess: object
    x_converged: bool
    x_tol: float
    err: float
    log: object = None  # TimeIterationLog
    trace: object = None  # {Nothing,IterationTrace}
    error_res: float = None

    @property
    def converged(self) -> bool:
        return bool(self.x_converged)

    @property
    def error_step(self) -> float:
        return self.err

    @property
    def error_eta(self) -> float:
        return self.err

    @property
    def error_eps(self) -> float:
        return self.error_res if self.error_res is not None else self.err