API Reference¶

Top-level API¶

The top-level namespace of dolo re-exports the model import and solution routines.

`PolicyLinearOperator` ¶

Bases: LinearOperator

Linear operator representing dF_2 = \partial F / \partial \varphi.

Acts on policy variations v defined on the continuous state grid

(L @ v)[n] = \sum_k F_{x, k}[n] \cdot v(S_k[n])

where v(S_k[n]) is the interpolation of v at the next-period state S_k[n].

Mirrors Dolo.jl's LL / LinearOperator struct.

Source code in src/dolo/numeric/residuals.py

class PolicyLinearOperator(LinearOperator):
    r"""Linear operator representing dF_2 = \partial F / \partial \varphi.

    Acts on policy variations v defined on the continuous state grid:
        (L @ v)[n] = \sum_k F_{x, k}[n] \cdot v(S_k[n])
    where v(S_k[n]) is the interpolation of v at the next-period state S_k[n].

    Mirrors Dolo.jl's LL / LinearOperator struct.
    """

    def __init__(
        self,
        dmodel: DiscretizedModel,
        F_x: Any,
        S: Any,
        engine: str = "jax",
        interp_method: str = "cubic",
        device: Optional[Any] = None,
        cr: Optional[CompiledResidual] = None,
    ):
        self.dmodel = dmodel
        self.F_x = F_x
        self.S = S
        self.engine = engine
        self.interp_method = interp_method
        self.device = device
        self.cr = cr

        self.n_regimes = dmodel.n_regimes
        self.endo_grid = dmodel.endo_grid
        self.N = self.endo_grid.nodes.shape[0]
        self.n_x = len(dmodel.model.calibration["controls"])
        self.n_s = len(dmodel.model.symbols["states"])
        self.grid_shape = tuple(int(x) for x in self.endo_grid.n)

        self.dim = self.n_regimes * self.N * self.n_x
        super().__init__(dtype=np.float64, shape=(self.dim, self.dim))

        if self.engine in ("jax", "gpu", "cpu") and self.cr is not None:
            import jax
            import jax.numpy as jnp
            from dolo.numeric.interpolation_jax import (
                prefilter_cubic_jax,
                eval_linear_point_jax,
                _eval_cubic_batch_jit,
            )

            a, b, n = self.cr.a, self.cr.b, self.cr.n
            grid_shape = self.grid_shape
            interp_method = self.interp_method

            def _apply_kernel(v_2d, F_x, S):
                v_reshaped = v_2d.reshape(grid_shape + (-1,))
                if interp_method == "cubic":
                    v_coeffs = prefilter_cubic_jax(grid_shape, v_reshaped)

                    def eval_branch(fx_k, s_k):
                        v_S_k = _eval_cubic_batch_jit(v_coeffs, s_k, a, b, n)
                        return jnp.einsum("nij,nj->ni", fx_k, v_S_k)
                else:
                    v_coeffs = v_reshaped

                    def eval_branch(fx_k, s_k):
                        v_S_k = jax.vmap(eval_linear_point_jax, in_axes=(None, 0, None, None, None))(
                            v_coeffs, s_k, a, b, n
                        )
                        return jnp.einsum("nij,nj->ni", fx_k, v_S_k)

                branch_results = jax.vmap(eval_branch)(F_x, S)
                return jnp.sum(branch_results, axis=0)

            self._apply_jit = jax.jit(_apply_kernel)
        else:
            self._apply_jit = None

    def _matvec(self, x: np.ndarray) -> np.ndarray:
        x_2d = x.reshape((self.N, self.n_x))
        out = self.apply(x_2d)
        return np.asarray(out).ravel()

    def apply(self, v: Union[GArray, np.ndarray, Any]) -> Union[GArray, np.ndarray, Any]:
        """Apply linear operator to a policy variation v on the grid."""
        if self.engine in ("jax", "gpu", "cpu"):
            import jax
            import jax.numpy as jnp
            from dolo.numeric.interpolation_jax import (
                prefilter_cubic_jax,
                eval_cubic_jax,
                eval_linear_jax,
            )

            is_garray = isinstance(v, GArray)
            v_data = v.data[0] if is_garray else v
            v_arr = jnp.asarray(v_data, dtype=jnp.float64)
            orig_ndim = v_arr.ndim
            if v_arr.ndim == 1:
                v_arr = v_arr.reshape((self.N, self.n_x))

            if self.device is not None:
                v_arr = jax.device_put(v_arr, self.device)

            if self._apply_jit is not None:
                out = self._apply_jit(v_arr, self.F_x, self.S)
            else:
                v_reshaped = v_arr.reshape(self.grid_shape + (-1,))
                if self.interp_method == "cubic":
                    v_coeffs = prefilter_cubic_jax(self.endo_grid, v_reshaped)
                    eval_fn = eval_cubic_jax
                else:
                    v_coeffs = v_reshaped
                    eval_fn = eval_linear_jax

                out = jnp.zeros_like(v_arr)
                n_branches = self.F_x.shape[0]
                for k in range(n_branches):
                    fx_k = self.F_x[k]
                    s_k = self.S[k]
                    v_S_k = eval_fn(self.endo_grid, v_coeffs, s_k)
                    out = out + jnp.einsum("nij,nj->ni", fx_k, v_S_k)

            if self.device is not None:
                out.block_until_ready()

            if is_garray:
                res = GArray.zeros_like(v)
                res[0] = np.asarray(out)
                return res
            elif isinstance(v, np.ndarray):
                if orig_ndim == 1:
                    return np.asarray(out).ravel()
                return np.asarray(out)
            return out

        else:
            # Vectorized NumPy evaluation
            is_garray = isinstance(v, GArray)
            v_arr = np.asarray(v.data if is_garray else v, dtype=np.float64)
            orig_ndim = v_arr.ndim
            if v_arr.ndim == 3 and v_arr.shape[0] == 1:
                v_2d = v_arr[0]
            elif v_arr.ndim == 1:
                v_2d = v_arr.reshape((self.N, self.n_x))
            else:
                v_2d = v_arr

            policy_ga = self.dmodel.initial_guess()
            if policy_ga.data.shape == v_2d.shape:
                policy_ga.data[...] = v_2d
            else:
                policy_ga[0] = v_2d
            dr_v = self.dmodel.decision_rule(policy_ga, interp_method=self.interp_method)

            out = np.zeros_like(v_2d)
            n_branches = self.F_x.shape[0]
            for k in range(n_branches):
                fx_k = np.asarray(self.F_x[k])
                s_k = np.asarray(self.S[k])
                v_S_k = dr_v.eval_ijs(0, k, s_k)
                out += np.einsum("nij,nj->ni", fx_k, v_S_k)

            if is_garray:
                res = GArray.zeros_like(v)
                res[0] = out
                return res
            elif orig_ndim == 1:
                return out.ravel()
            return out

    def __call__(self, v: Any) -> Any:
        return self.apply(v)

    def __matmul__(self, v: Any) -> Any:
        return self.apply(v)

    def __mul__(self, alpha: Union[float, int]) -> PolicyLinearOperator:
        return PolicyLinearOperator(
            self.dmodel,
            self.F_x * alpha,
            self.S,
            engine=self.engine,
            interp_method=self.interp_method,
            device=self.device,
            cr=self.cr,
        )

    def __rmul__(self, alpha: Union[float, int]) -> PolicyLinearOperator:
        return self.__mul__(alpha)

    def preconditioned(self, J: Any) -> PolicyLinearOperator:
        """Return preconditioned operator J^{-1} L, where J = dF_1 is block-diagonal."""
        n_branches = self.F_x.shape[0]
        if self.engine in ("jax", "gpu", "cpu"):
            import jax
            import jax.numpy as jnp

            J_arr = jnp.asarray(J.data[0] if isinstance(J, GArray) else J, dtype=jnp.float64)
            if self.device is not None:
                J_arr = jax.device_put(J_arr, self.device)
            new_Fx = jnp.stack(
                [jnp.linalg.solve(J_arr, self.F_x[k]) for k in range(n_branches)],
                axis=0,
            )
        else:
            J_arr = np.asarray(J.data[0] if isinstance(J, GArray) else J, dtype=np.float64)
            new_Fx = np.stack(
                [np.linalg.solve(J_arr, np.asarray(self.F_x[k])) for k in range(n_branches)],
                axis=0,
            )

        return PolicyLinearOperator(
            self.dmodel,
            new_Fx,
            self.S,
            engine=self.engine,
            interp_method=self.interp_method,
            device=self.device,
            cr=self.cr,
        )

    def as_matrix(self) -> np.ndarray:
        """Compute explicit (dim, dim) dense matrix representation."""
        dim = self.dim
        mat = np.zeros((dim, dim))
        e = np.zeros(dim)
        for i in range(dim):
            if i > 0:
                e[i - 1] = 0.0
            e[i] = 1.0
            mat[:, i] = self._matvec(e)
        return mat

    def solve(
        self,
        b: Any,
        linsolve: str = "neumann",
        tol: float = 1e-8,
        maxit: int = 50,
        verbose: bool = False,
    ) -> Any:
        """Solve (I - L) x = b using Neumann series or GMRES."""
        if linsolve == "neumann":
            x = b.copy()
            du = b.copy()
            for it in range(maxit):
                du = self.apply(du)
                x = x + du
                err = np.abs(np.asarray(du)).max()
                if err < tol:
                    break
            return x
        elif linsolve == "gmres":
            from scipy.sparse.linalg import gmres as scipy_gmres

            I_minus_L = LinearOperator(
                dtype=np.float64,
                shape=self.shape,
                matvec=lambda v: v - self._matvec(v),
            )
            b_np = np.asarray(b.data if isinstance(b, GArray) else b)
            b_flat = b_np.ravel()
            sol_flat, info = scipy_gmres(I_minus_L, b_flat, rtol=tol, atol=0.0, maxiter=maxit)
            sol_reshaped = sol_flat.reshape(b_np.shape)
            if isinstance(b, GArray):
                out = GArray.zeros_like(b)
                out.data[...] = sol_reshaped
                return out
            return sol_reshaped
        else:
            raise ValueError(f"Unknown linear solver: '{linsolve}'. Use 'neumann' or 'gmres'.")

`apply(v)` ¶

Apply linear operator to a policy variation v on the grid.

Source code in src/dolo/numeric/residuals.py

def apply(self, v: Union[GArray, np.ndarray, Any]) -> Union[GArray, np.ndarray, Any]:
    """Apply linear operator to a policy variation v on the grid."""
    if self.engine in ("jax", "gpu", "cpu"):
        import jax
        import jax.numpy as jnp
        from dolo.numeric.interpolation_jax import (
            prefilter_cubic_jax,
            eval_cubic_jax,
            eval_linear_jax,
        )

        is_garray = isinstance(v, GArray)
        v_data = v.data[0] if is_garray else v
        v_arr = jnp.asarray(v_data, dtype=jnp.float64)
        orig_ndim = v_arr.ndim
        if v_arr.ndim == 1:
            v_arr = v_arr.reshape((self.N, self.n_x))

        if self.device is not None:
            v_arr = jax.device_put(v_arr, self.device)

        if self._apply_jit is not None:
            out = self._apply_jit(v_arr, self.F_x, self.S)
        else:
            v_reshaped = v_arr.reshape(self.grid_shape + (-1,))
            if self.interp_method == "cubic":
                v_coeffs = prefilter_cubic_jax(self.endo_grid, v_reshaped)
                eval_fn = eval_cubic_jax
            else:
                v_coeffs = v_reshaped
                eval_fn = eval_linear_jax

            out = jnp.zeros_like(v_arr)
            n_branches = self.F_x.shape[0]
            for k in range(n_branches):
                fx_k = self.F_x[k]
                s_k = self.S[k]
                v_S_k = eval_fn(self.endo_grid, v_coeffs, s_k)
                out = out + jnp.einsum("nij,nj->ni", fx_k, v_S_k)

        if self.device is not None:
            out.block_until_ready()

        if is_garray:
            res = GArray.zeros_like(v)
            res[0] = np.asarray(out)
            return res
        elif isinstance(v, np.ndarray):
            if orig_ndim == 1:
                return np.asarray(out).ravel()
            return np.asarray(out)
        return out

    else:
        # Vectorized NumPy evaluation
        is_garray = isinstance(v, GArray)
        v_arr = np.asarray(v.data if is_garray else v, dtype=np.float64)
        orig_ndim = v_arr.ndim
        if v_arr.ndim == 3 and v_arr.shape[0] == 1:
            v_2d = v_arr[0]
        elif v_arr.ndim == 1:
            v_2d = v_arr.reshape((self.N, self.n_x))
        else:
            v_2d = v_arr

        policy_ga = self.dmodel.initial_guess()
        if policy_ga.data.shape == v_2d.shape:
            policy_ga.data[...] = v_2d
        else:
            policy_ga[0] = v_2d
        dr_v = self.dmodel.decision_rule(policy_ga, interp_method=self.interp_method)

        out = np.zeros_like(v_2d)
        n_branches = self.F_x.shape[0]
        for k in range(n_branches):
            fx_k = np.asarray(self.F_x[k])
            s_k = np.asarray(self.S[k])
            v_S_k = dr_v.eval_ijs(0, k, s_k)
            out += np.einsum("nij,nj->ni", fx_k, v_S_k)

        if is_garray:
            res = GArray.zeros_like(v)
            res[0] = out
            return res
        elif orig_ndim == 1:
            return out.ravel()
        return out

`as_matrix()` ¶

Compute explicit (dim, dim) dense matrix representation.

Source code in src/dolo/numeric/residuals.py

def as_matrix(self) -> np.ndarray:
    """Compute explicit (dim, dim) dense matrix representation."""
    dim = self.dim
    mat = np.zeros((dim, dim))
    e = np.zeros(dim)
    for i in range(dim):
        if i > 0:
            e[i - 1] = 0.0
        e[i] = 1.0
        mat[:, i] = self._matvec(e)
    return mat

`preconditioned(J)` ¶

Return preconditioned operator J^{-1} L, where J = dF_1 is block-diagonal.

Source code in src/dolo/numeric/residuals.py

def preconditioned(self, J: Any) -> PolicyLinearOperator:
    """Return preconditioned operator J^{-1} L, where J = dF_1 is block-diagonal."""
    n_branches = self.F_x.shape[0]
    if self.engine in ("jax", "gpu", "cpu"):
        import jax
        import jax.numpy as jnp

        J_arr = jnp.asarray(J.data[0] if isinstance(J, GArray) else J, dtype=jnp.float64)
        if self.device is not None:
            J_arr = jax.device_put(J_arr, self.device)
        new_Fx = jnp.stack(
            [jnp.linalg.solve(J_arr, self.F_x[k]) for k in range(n_branches)],
            axis=0,
        )
    else:
        J_arr = np.asarray(J.data[0] if isinstance(J, GArray) else J, dtype=np.float64)
        new_Fx = np.stack(
            [np.linalg.solve(J_arr, np.asarray(self.F_x[k])) for k in range(n_branches)],
            axis=0,
        )

    return PolicyLinearOperator(
        self.dmodel,
        new_Fx,
        self.S,
        engine=self.engine,
        interp_method=self.interp_method,
        device=self.device,
        cr=self.cr,
    )

`solve(b, linsolve='neumann', tol=1e-08, maxit=50, verbose=False)` ¶

Solve (I - L) x = b using Neumann series or GMRES.

Source code in src/dolo/numeric/residuals.py

def solve(
    self,
    b: Any,
    linsolve: str = "neumann",
    tol: float = 1e-8,
    maxit: int = 50,
    verbose: bool = False,
) -> Any:
    """Solve (I - L) x = b using Neumann series or GMRES."""
    if linsolve == "neumann":
        x = b.copy()
        du = b.copy()
        for it in range(maxit):
            du = self.apply(du)
            x = x + du
            err = np.abs(np.asarray(du)).max()
            if err < tol:
                break
        return x
    elif linsolve == "gmres":
        from scipy.sparse.linalg import gmres as scipy_gmres

        I_minus_L = LinearOperator(
            dtype=np.float64,
            shape=self.shape,
            matvec=lambda v: v - self._matvec(v),
        )
        b_np = np.asarray(b.data if isinstance(b, GArray) else b)
        b_flat = b_np.ravel()
        sol_flat, info = scipy_gmres(I_minus_L, b_flat, rtol=tol, atol=0.0, maxiter=maxit)
        sol_reshaped = sol_flat.reshape(b_np.shape)
        if isinstance(b, GArray):
            out = GArray.zeros_like(b)
            out.data[...] = sol_reshaped
            return out
        return sol_reshaped
    else:
        raise ValueError(f"Unknown linear solver: '{linsolve}'. Use 'neumann' or 'gmres'.")

`F(model, *args, engine='auto', backend=None, interp_method='cubic', device=None, **kwargs)` ¶

Evaluate dynamic Euler equation residuals across grid points or at a single point.

Usage¶

Grid-level evaluation (mirroring Julia Dolo.jl): F(model, controls, dr=None, engine="auto") F(model, x0, x0) # dr initialized to initial guess x0 F(model, x0) # dr defaults to x0
Point-level evaluation: F(model, s, x, dr)

Parameters¶

model : Model, ModelV2, or DiscretizedModel The economic model. controls : GArray, np.ndarray, or jnp.ndarray Control policy values across the grid (or point state s if point-call). dr : DecisionRule, GArray, or array, optional Decision rule (policy) to evaluate next-period controls φ(S). If omitted or equal to controls, fits φ on controls. engine / backend : str ("auto", "jax", "gpu", "cpu", "vectorized") Computational backend to execute the residual evaluation. interp_method : str ("cubic" or "linear") Interpolation method for decision rule φ. device : jax.Device, optional Target JAX device (e.g. jax.devices('gpu')[0]).

Source code in src/dolo/numeric/residuals.py

def F(
    model: Any,
    *args,
    engine: str = "auto",
    backend: Optional[str] = None,
    interp_method: str = "cubic",
    device: Optional[Any] = None,
    **kwargs,
) -> Union[GArray, np.ndarray, Any]:
    """Evaluate dynamic Euler equation residuals across grid points or at a single point.

    Usage
    -----
    1. Grid-level evaluation (mirroring Julia Dolo.jl):
       F(model, controls, dr=None, engine="auto")
       F(model, x0, x0)  # dr initialized to initial guess x0
       F(model, x0)      # dr defaults to x0

    2. Point-level evaluation:
       F(model, s, x, dr)

    Parameters
    ----------
    model : Model, ModelV2, or DiscretizedModel
        The economic model.
    controls : GArray, np.ndarray, or jnp.ndarray
        Control policy values across the grid (or point state s if point-call).
    dr : DecisionRule, GArray, or array, optional
        Decision rule (policy) to evaluate next-period controls φ(S).
        If omitted or equal to controls, fits φ on `controls`.
    engine / backend : str ("auto", "jax", "gpu", "cpu", "vectorized")
        Computational backend to execute the residual evaluation.
    interp_method : str ("cubic" or "linear")
        Interpolation method for decision rule φ.
    device : jax.Device, optional
        Target JAX device (e.g. jax.devices('gpu')[0]).
    """
    if len(args) == 0:
        raise ValueError("F requires at least 2 arguments: (model, controls) or (model, s, x, dr)")

    # 1. Point-level dispatch: F(model, s, x, dr)
    if len(args) >= 2 and (
        isinstance(args[0], (QP, tuple))
        or (isinstance(args[0], np.ndarray) and args[0].ndim == 1)
    ):
        s, x = args[0], args[1]
        dr = args[2] if len(args) > 2 else None
        dmodel = model if isinstance(model, DiscretizedModel) else DiscretizedModel(model, **kwargs)
        eng = _resolve_engine(dmodel, engine, backend)
        if eng in ("jax", "gpu", "cpu"):
            cr = _get_or_compile_jax_residual(dmodel, interp_method=interp_method)
            coeffs = cr.prefilter(dr if dr is not None else x)
            return cr.F_point(s if not isinstance(s, QP) else s.val, x, coeffs)
        else:
            # Vectorized / point transition
            dr_obj = dr if isinstance(dr, DecisionRule) else dmodel.decision_rule(dr, interp_method=interp_method)
            s_qp = s if isinstance(s, QP) else QP(0, s)
            res = np.zeros_like(x)
            n_m = len(dmodel.model.symbols["exogenous"]) if "exogenous" in dmodel.model.symbols else len(dmodel.model.calibration.get("exogenous", []))
            for prob, S in dmodel.tau.point_transitions(s_qp, x):
                rr = dmodel.arbitrage(s_qp.val[:n_m], s_qp.val[n_m:], x, S.val[:n_m], S.val[n_m:], dr_obj(S))
                res += prob * rr
            return res

    # 2. Grid-level dispatch: F(model, controls, dr=None)
    controls = args[0]
    dr = args[1] if len(args) > 1 else None

    # Resolve discretized model
    if isinstance(model, DiscretizedModel):
        dmodel = model
    else:
        grid_options = kwargs.pop("grid_options", None)
        dprocess_options = kwargs.pop("dprocess_options", {})
        dmodel = DiscretizedModel(model, grid_options=grid_options, dprocess_options=dprocess_options, **kwargs)

    eng = _resolve_engine(dmodel, engine, backend)

    if eng in ("jax", "gpu", "cpu"):
        import jax
        import jax.numpy as jnp

        target_dev = device
        if target_dev is 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=interp_method)
        policy_for_coeffs = dr if dr is not None else controls
        coeffs = cr.prefilter(policy_for_coeffs)

        s_nodes = jnp.asarray(dmodel.endo_nodes, dtype=jnp.float64)
        if isinstance(controls, GArray):
            x_vals = jnp.asarray(controls[0] if dmodel.n_regimes == 1 else controls.data, dtype=jnp.float64)
        else:
            c_arr = jnp.asarray(controls, dtype=jnp.float64)
            x_vals = c_arr[0] if (dmodel.n_regimes == 1 and c_arr.ndim == 3) else c_arr

        if target_dev is not None:
            s_nodes = jax.device_put(s_nodes, target_dev)
            x_vals = jax.device_put(x_vals, target_dev)
            coeffs = jax.device_put(coeffs, target_dev)

        if dmodel.n_regimes == 1:
            res_jax = cr.F_vmap(s_nodes, x_vals, coeffs)
            if target_dev is not None:
                res_jax.block_until_ready()
            if isinstance(controls, GArray):
                out = GArray.zeros_like(controls)
                out[0] = np.asarray(res_jax)
                return out
            return res_jax
        else:
            return cr.F(controls, coeffs)

    elif eng == "vectorized":
        if isinstance(dr, DecisionRule):
            dr_obj = dr
        elif dr is None or dr is controls or isinstance(dr, (GArray, np.ndarray)):
            policy_val = controls if dr is None or dr is controls else dr
            if not isinstance(policy_val, GArray):
                policy_ga = dmodel.initial_guess()
                if policy_ga.data.shape == policy_val.shape:
                    policy_ga.data[...] = policy_val
                elif policy_ga[0].shape == policy_val.shape:
                    policy_ga[0] = policy_val
                policy_val = policy_ga
            dr_obj = dmodel.decision_rule(policy_val, interp_method=interp_method)
        else:
            dr_obj = dr

        controls_ga = controls if isinstance(controls, GArray) else dmodel.initial_guess()
        if not isinstance(controls, GArray):
            if controls_ga.data.shape == controls.shape:
                controls_ga.data[...] = controls
            elif controls_ga[0].shape == controls.shape:
                controls_ga[0] = controls

        return euler_residuals_tau(dmodel, controls_ga, dr_obj)

    raise ValueError(f"Unsupported engine: {eng}")

`F_and_dF_1(model, *args, engine='auto', backend=None, interp_method='cubic', diff_mode='fwd', device=None, **kwargs)` ¶

Simultaneously evaluate residuals and Jacobian w.r.t controls in a single pass.

Returns¶

(res, dres) : Tuple of residuals and Jacobian tensor

Source code in src/dolo/numeric/residuals.py

def F_and_dF_1(
    model: Any,
    *args,
    engine: str = "auto",
    backend: Optional[str] = None,
    interp_method: str = "cubic",
    diff_mode: str = "fwd",
    device: Optional[Any] = None,
    **kwargs,
) -> Tuple[Any, Any]:
    """Simultaneously evaluate residuals and Jacobian w.r.t controls in a single pass.

    Returns
    -------
    (res, dres) : Tuple of residuals and Jacobian tensor
    """
    if len(args) == 0:
        raise ValueError("F_and_dF_1 requires at least 1 positional argument: controls")

    controls = args[0]
    dr = args[1] if len(args) > 1 else None

    if isinstance(model, DiscretizedModel):
        dmodel = model
    else:
        grid_options = kwargs.pop("grid_options", None)
        dprocess_options = kwargs.pop("dprocess_options", {})
        dmodel = DiscretizedModel(model, grid_options=grid_options, dprocess_options=dprocess_options, **kwargs)

    eng = _resolve_engine(dmodel, engine, backend)

    if eng in ("jax", "gpu", "cpu"):
        import jax
        import jax.numpy as jnp

        target_dev = device
        if target_dev is 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=interp_method, diff_mode=diff_mode)
        policy_for_coeffs = dr if dr is not None else controls
        coeffs = cr.prefilter(policy_for_coeffs)

        s_nodes = jnp.asarray(dmodel.endo_nodes, dtype=jnp.float64)
        if isinstance(controls, GArray):
            x_vals = jnp.asarray(controls[0] if dmodel.n_regimes == 1 else controls.data, dtype=jnp.float64)
        else:
            c_arr = jnp.asarray(controls, dtype=jnp.float64)
            x_vals = c_arr[0] if (dmodel.n_regimes == 1 and c_arr.ndim == 3) else c_arr

        if target_dev is not None:
            s_nodes = jax.device_put(s_nodes, target_dev)
            x_vals = jax.device_put(x_vals, target_dev)
            coeffs = jax.device_put(coeffs, target_dev)

        if dmodel.n_regimes == 1:
            res, dres = cr.F_and_dF_vmap(s_nodes, x_vals, coeffs)
            if target_dev is not None:
                res.block_until_ready()
                dres.block_until_ready()
            if isinstance(controls, GArray):
                out = GArray.zeros_like(controls)
                out[0] = np.asarray(res)
                return out, dres
            return res, dres
        else:
            return cr.F(controls, coeffs), cr.dF(controls, coeffs)

    elif eng == "vectorized":
        res = F(dmodel, controls, dr, engine="vectorized", interp_method=interp_method, **kwargs)
        dres = dF_1(dmodel, controls, dr, engine="vectorized", interp_method=interp_method, **kwargs)
        return res, dres

    raise ValueError(f"Unsupported engine: {eng}")

`colored(text, color=None, on_color=None, attrs=None)` ¶

Colorize text.

Available text colors

red, green, yellow, blue, magenta, cyan, white.

Available text highlights

on_red, on_green, on_yellow, on_blue, on_magenta, on_cyan, on_white.

Available attributes

bold, dark, underline, blink, reverse, concealed.

Example

colored('Hello, World!', 'red', 'on_grey', ['blue', 'blink']) colored('Hello, World!', 'green')

Source code in src/dolo/misc/termcolor.py

def colored(text, color=None, on_color=None, attrs=None):
    """Colorize text.

    Available text colors:
        red, green, yellow, blue, magenta, cyan, white.

    Available text highlights:
        on_red, on_green, on_yellow, on_blue, on_magenta, on_cyan, on_white.

    Available attributes:
        bold, dark, underline, blink, reverse, concealed.

    Example:
        colored('Hello, World!', 'red', 'on_grey', ['blue', 'blink'])
        colored('Hello, World!', 'green')
    """
    if os.getenv("ANSI_COLORS_DISABLED") is None:
        fmt_str = "\033[%dm%s"
        if color is not None:
            text = fmt_str % (COLORS[color], text)

        if on_color is not None:
            text = fmt_str % (HIGHLIGHTS[on_color], text)

        if attrs is not None:
            for attr in attrs:
                text = fmt_str % (ATTRIBUTES[attr], text)

        text += RESET
    return text

`dF_1(model, *args, engine='auto', backend=None, interp_method='cubic', diff_mode='fwd', device=None, **kwargs)` ¶

Evaluate Jacobian of dynamic Euler residuals w.r.t controls across grid points.

Mirrors Julia Dolo.jl: dF_1(model, controls, φ=None)

Parameters¶

model : Model, ModelV2, or DiscretizedModel The economic model. controls : GArray, np.ndarray, or jnp.ndarray Control policy values across the grid. dr : DecisionRule, GArray, or array, optional Decision rule (policy). Defaults to controls if omitted. engine / backend : str ("auto", "jax", "gpu", "cpu", "vectorized") Backend engine. JAX computes exact analytical autodiff derivatives. interp_method : str ("cubic" or "linear") Interpolation method. diff_mode : str ("fwd" or "rev") Autodiff mode for JAX ('fwd' for jacfwd, 'rev' for jacrev).

Source code in src/dolo/numeric/residuals.py

def dF_1(
    model: Any,
    *args,
    engine: str = "auto",
    backend: Optional[str] = None,
    interp_method: str = "cubic",
    diff_mode: str = "fwd",
    device: Optional[Any] = None,
    **kwargs,
) -> Any:
    """Evaluate Jacobian of dynamic Euler residuals w.r.t controls across grid points.

    Mirrors Julia Dolo.jl: dF_1(model, controls, φ=None)

    Parameters
    ----------
    model : Model, ModelV2, or DiscretizedModel
        The economic model.
    controls : GArray, np.ndarray, or jnp.ndarray
        Control policy values across the grid.
    dr : DecisionRule, GArray, or array, optional
        Decision rule (policy). Defaults to `controls` if omitted.
    engine / backend : str ("auto", "jax", "gpu", "cpu", "vectorized")
        Backend engine. JAX computes exact analytical autodiff derivatives.
    interp_method : str ("cubic" or "linear")
        Interpolation method.
    diff_mode : str ("fwd" or "rev")
        Autodiff mode for JAX ('fwd' for jacfwd, 'rev' for jacrev).
    """
    if len(args) == 0:
        raise ValueError("dF_1 requires at least 1 positional argument: controls")

    controls = args[0]
    dr = args[1] if len(args) > 1 else None

    if isinstance(model, DiscretizedModel):
        dmodel = model
    else:
        grid_options = kwargs.pop("grid_options", None)
        dprocess_options = kwargs.pop("dprocess_options", {})
        dmodel = DiscretizedModel(model, grid_options=grid_options, dprocess_options=dprocess_options, **kwargs)

    eng = _resolve_engine(dmodel, engine, backend)

    if eng in ("jax", "gpu", "cpu"):
        import jax
        import jax.numpy as jnp

        target_dev = device
        if target_dev is 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=interp_method, diff_mode=diff_mode)
        policy_for_coeffs = dr if dr is not None else controls
        coeffs = cr.prefilter(policy_for_coeffs)

        s_nodes = jnp.asarray(dmodel.endo_nodes, dtype=jnp.float64)
        if isinstance(controls, GArray):
            x_vals = jnp.asarray(controls[0] if dmodel.n_regimes == 1 else controls.data, dtype=jnp.float64)
        else:
            c_arr = jnp.asarray(controls, dtype=jnp.float64)
            x_vals = c_arr[0] if (dmodel.n_regimes == 1 and c_arr.ndim == 3) else c_arr

        if target_dev is not None:
            s_nodes = jax.device_put(s_nodes, target_dev)
            x_vals = jax.device_put(x_vals, target_dev)
            coeffs = jax.device_put(coeffs, target_dev)

        if dmodel.n_regimes == 1:
            jac = cr.dF_vmap(s_nodes, x_vals, coeffs)
            if target_dev is not None:
                jac.block_until_ready()
            return jac
        else:
            return cr.dF(controls, coeffs)

    elif eng == "vectorized":
        # Finite-difference derivative using SerialDifferentiableFunction
        controls_ga = controls if isinstance(controls, GArray) else dmodel.initial_guess()
        if not isinstance(controls, GArray):
            if controls_ga.data.shape == controls.shape:
                controls_ga.data[...] = controls
            elif controls_ga[0].shape == controls.shape:
                controls_ga[0] = controls

        if isinstance(dr, DecisionRule):
            dr_obj = dr
        else:
            dr_obj = dmodel.decision_rule(controls_ga if dr is None else dr, interp_method=interp_method)

        n_x = len(dmodel.model.calibration["controls"])
        fn = lambda x_val: euler_residuals_tau(dmodel, controls_ga.unravel(x_val), dr_obj).ravel().reshape((-1, n_x))
        dfn = SerialDifferentiableFunction(fn)
        _, jacs = dfn(controls_ga.ravel().reshape((-1, n_x)))
        return jacs

    raise ValueError(f"Unsupported engine: {eng}")

`dF_2(model, *args, engine='auto', backend=None, interp_method='cubic', device=None, **kwargs)` ¶

Evaluate derivative of dynamic Euler residuals w.r.t the decision rule φ.

Mirrors Julia Dolo.jl: dF_2(model, s, x, φ) and dF_2(model, controls, φ=None)

Usage¶

Grid-level evaluation: L = dF_2(model, controls, dr=None, engine="auto") L @ v # apply operator to variation v L.as_matrix() # explicit matrix
Point-level evaluation: branches = dF_2(model, s, x, dr) # returns tuple of (F_x, S) for each shock branch k

Parameters¶

model : Model, ModelV2, or DiscretizedModel The economic model. controls : GArray, np.ndarray, or jnp.ndarray Control policy values across the grid (or point state s if point call). dr : DecisionRule, GArray, or array, optional Decision rule (policy). Defaults to controls if omitted. engine / backend : str ("auto", "jax", "gpu", "cpu", "vectorized") Backend engine. interp_method : str ("cubic" or "linear") Interpolation method for decision rule φ. device : jax.Device, optional Target JAX device (e.g. jax.devices('gpu')[0]).

Source code in src/dolo/numeric/residuals.py

def dF_2(
    model: Any,
    *args,
    engine: str = "auto",
    backend: Optional[str] = None,
    interp_method: str = "cubic",
    device: Optional[Any] = None,
    **kwargs,
) -> Union[Tuple[Tuple[Any, Any], ...], PolicyLinearOperator]:
    r"""Evaluate derivative of dynamic Euler residuals w.r.t the decision rule φ.

    Mirrors Julia Dolo.jl: dF_2(model, s, x, φ) and dF_2(model, controls, φ=None)

    Usage
    -----
    1. Grid-level evaluation:
       L = dF_2(model, controls, dr=None, engine="auto")
       L @ v        # apply operator to variation v
       L.as_matrix() # explicit matrix

    2. Point-level evaluation:
       branches = dF_2(model, s, x, dr)
       # returns tuple of (F_x, S) for each shock branch k

    Parameters
    ----------
    model : Model, ModelV2, or DiscretizedModel
        The economic model.
    controls : GArray, np.ndarray, or jnp.ndarray
        Control policy values across the grid (or point state s if point call).
    dr : DecisionRule, GArray, or array, optional
        Decision rule (policy). Defaults to `controls` if omitted.
    engine / backend : str ("auto", "jax", "gpu", "cpu", "vectorized")
        Backend engine.
    interp_method : str ("cubic" or "linear")
        Interpolation method for decision rule φ.
    device : jax.Device, optional
        Target JAX device (e.g. jax.devices('gpu')[0]).
    """
    if len(args) == 0:
        raise ValueError("dF_2 requires at least 1 argument: (model, controls) or (model, s, x, dr)")

    # 1. Point-level dispatch: dF_2(model, s, x, dr)
    if len(args) >= 2 and (
        isinstance(args[0], (QP, tuple))
        or (isinstance(args[0], np.ndarray) and args[0].ndim == 1)
    ):
        s, x = args[0], args[1]
        dr = args[2] if len(args) > 2 else None
        dmodel = model if isinstance(model, DiscretizedModel) else DiscretizedModel(model, **kwargs)
        eng = _resolve_engine(dmodel, engine, backend)

        if eng in ("jax", "gpu", "cpu"):
            import jax
            import jax.numpy as jnp
            from dolo.numeric.interpolation_jax import eval_cubic_point_jax, eval_linear_point_jax

            cr = _get_or_compile_jax_residual(dmodel, interp_method=interp_method)
            policy_for_coeffs = dr if dr is not None else x
            coeffs = cr.prefilter(policy_for_coeffs)

            s_pt = jnp.asarray(s if not isinstance(s, QP) else s.val, dtype=jnp.float64)
            x_pt = jnp.asarray(x, dtype=jnp.float64)
            m_val, inodes, weights = cr.regime_data[0]
            n_inodes = len(weights)

            branches = []
            for k in range(n_inodes):
                M = inodes[k]
                w = weights[k]
                S = cr.g_pt(m_val, s_pt, x_pt, M, cr.parms)
                if interp_method == "cubic":
                    X = eval_cubic_point_jax(coeffs, S, cr.a, cr.b, cr.n)
                else:
                    X = eval_linear_point_jax(coeffs, S, cr.a, cr.b, cr.n)
                df_dX = jax.jacobian(lambda u: cr.f_pt(m_val, s_pt, x_pt, M, S, u, cr.parms))(X)
                branches.append((np.asarray(w * df_dX), np.asarray(S)))
            return tuple(branches)
        else:
            # Vectorized point
            dr_obj = dr if isinstance(dr, DecisionRule) else dmodel.decision_rule(dr if dr is not None else x, interp_method=interp_method)
            s_val = s.val if isinstance(s, QP) else s
            n_m = len(dmodel.model.symbols["exogenous"]) if "exogenous" in dmodel.model.symbols else len(dmodel.model.calibration.get("exogenous", []))
            if isinstance(s, QP):
                m_val = s_val[:n_m]
                s_endo = s_val[n_m:]
            else:
                m_val = np.asarray(dmodel.dprocess.node(0))
                s_endo = s_val

            n_inodes = dmodel.dprocess.n_inodes(0)

            branches = []
            for k in range(n_inodes):
                M_val = dmodel.dprocess.inode(0, k)
                w = float(dmodel.dprocess.iweight(0, k))
                S_endo = dmodel.g(m_val[None, :], s_endo[None, :], x[None, :], M_val[None, :], dmodel.parms)[0]
                val_next = np.concatenate([M_val, S_endo])
                target_regime = 0 if dmodel.dprocess.n_nodes == 1 else k
                S_qp = QP(loc=(target_regime, S_endo), val=val_next)
                X_next = dr_obj(S_qp)

                ff = SerialDifferentiableFunction(
                    lambda u: dmodel.arbitrage(m_val[None, :], s_endo[None, :], x[None, :], M_val[None, :], S_endo[None, :], u)
                )
                _, rr_X = ff(X_next[None, :])
                branches.append((w * rr_X[0], S_endo))
            return tuple(branches)

    # 2. Grid-level dispatch: dF_2(model, controls, dr=None)
    controls = args[0]
    dr = args[1] if len(args) > 1 else None

    if isinstance(model, DiscretizedModel):
        dmodel = model
    else:
        grid_options = kwargs.pop("grid_options", None)
        dprocess_options = kwargs.pop("dprocess_options", {})
        dmodel = DiscretizedModel(model, grid_options=grid_options, dprocess_options=dprocess_options, **kwargs)

    eng = _resolve_engine(dmodel, engine, backend)

    if eng in ("jax", "gpu", "cpu"):
        import jax
        import jax.numpy as jnp
        from dolo.numeric.interpolation_jax import eval_cubic_point_jax, eval_linear_point_jax

        target_dev = device
        if target_dev is 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=interp_method)
        policy_for_coeffs = dr if dr is not None else controls
        coeffs = cr.prefilter(policy_for_coeffs)

        s_nodes = jnp.asarray(dmodel.endo_nodes, dtype=jnp.float64)
        if isinstance(controls, GArray):
            x_vals = jnp.asarray(controls[0] if dmodel.n_regimes == 1 else controls.data, dtype=jnp.float64)
        else:
            c_arr = jnp.asarray(controls, dtype=jnp.float64)
            x_vals = c_arr[0] if (dmodel.n_regimes == 1 and c_arr.ndim == 3) else c_arr

        if target_dev is not None:
            s_nodes = jax.device_put(s_nodes, target_dev)
            x_vals = jax.device_put(x_vals, target_dev)
            coeffs = jax.device_put(coeffs, target_dev)

        f_pt = cr.f_pt
        g_pt = cr.g_pt
        parms = cr.parms
        m_val, inodes, weights = cr.regime_data[0]
        n_inodes = len(weights)
        a, b, n = cr.a, cr.b, cr.n

        def branch_eval(s, x, coeffs, M, w):
            S = g_pt(m_val, s, x, M, parms)
            if interp_method == "cubic":
                X = eval_cubic_point_jax(coeffs, S, a, b, n)
            else:
                X = eval_linear_point_jax(coeffs, S, a, b, n)
            df_dX = jax.jacobian(lambda u: f_pt(m_val, s, x, M, S, u, parms))(X)
            return w * df_dX, S

        branch_vmap = jax.jit(jax.vmap(branch_eval, in_axes=(0, 0, None, None, None)))

        F_x_list = []
        S_list = []
        for k in range(n_inodes):
            M = inodes[k]
            w = weights[k]
            fx_k, s_k = branch_vmap(s_nodes, x_vals, coeffs, M, w)
            F_x_list.append(fx_k)
            S_list.append(s_k)

        F_x_all = jnp.stack(F_x_list, axis=0)  # (K, N, n_x, n_x)
        S_all = jnp.stack(S_list, axis=0)      # (K, N, n_s)

        if target_dev is not None:
            F_x_all.block_until_ready()
            S_all.block_until_ready()

        return PolicyLinearOperator(
            dmodel,
            F_x_all,
            S_all,
            engine=eng,
            interp_method=interp_method,
            device=target_dev,
            cr=cr,
        )

    elif eng == "vectorized":
        # Pure NumPy evaluation
        controls_ga = controls if isinstance(controls, GArray) else dmodel.initial_guess()
        if not isinstance(controls, GArray):
            if controls_ga.data.shape == controls.shape:
                controls_ga.data[...] = controls
            elif controls_ga[0].shape == controls.shape:
                controls_ga[0] = controls

        if isinstance(dr, DecisionRule):
            dr_obj = dr
        else:
            dr_obj = dmodel.decision_rule(controls_ga if dr is None else dr, interp_method=interp_method)

        s_nodes = dmodel.endo_nodes
        x_vals = controls_ga[0] if isinstance(controls_ga, GArray) else (controls_ga[0] if controls_ga.ndim == 3 else controls_ga)
        m_val = dmodel.dprocess.node(0)
        n_inodes = dmodel.dprocess.n_inodes(0)
        N = s_nodes.shape[0]
        n_x = len(dmodel.model.calibration["controls"])
        m_tiled = np.tile(m_val, (N, 1))

        F_x_list = []
        S_list = []
        for k in range(n_inodes):
            M_val = dmodel.dprocess.inode(0, k)
            M_tiled = np.tile(M_val, (N, 1))
            w = float(dmodel.dprocess.iweight(0, k))

            S = dmodel.g(m_tiled, s_nodes, x_vals, M_tiled, dmodel.parms)
            X = dr_obj.eval_ijs(0, k, S)

            ff = SerialDifferentiableFunction(
                lambda u: dmodel.arbitrage(m_val, s_nodes, x_vals, M_val, S, u)
            )
            _, rr_X = ff(X)
            F_x_list.append(w * rr_X)
            S_list.append(S)

        F_x_all = np.stack(F_x_list, axis=0)  # (K, N, n_x, n_x)
        S_all = np.stack(S_list, axis=0)      # (K, N, n_s)

        return PolicyLinearOperator(
            dmodel,
            F_x_all,
            S_all,
            engine="vectorized",
            interp_method=interp_method,
            device=None,
            cr=None,
        )

    raise ValueError(f"Unsupported engine: {eng}")

`deterministic_solve(model, *, verbose=True, ignore_constraints=False, exogenous=None, s0=None, m0=None, T=100, maxit=100, initial_guess=None, solver='ncpsolve', keep_steady_state=False, s1=None, shocks=None, tol=1e-06)` ¶

Computes a perfect foresight simulation using a stacked-time algorithm.

Typical simulation exercises are: - start from an out-of-equilibrium exogenous and/or endogenous state: specify s0 and or m0. Missing values are taken from the calibration (model.calibration). - specify an exogenous path for shocks exogenous. Initial exogenous state m0 is then first value of exogenous values. Economy is supposed to have been at the equilibrium for \(t<0\), which pins down initial endogenous state s0. x0 is a jump variable.

If \(s0\) is not specified it is then set equal to the steady-state consistent with the first value

The initial state is specified either by providing a series of exogenous shocks and assuming the model is initially in equilibrium with the first value of the shock, or by specifying an initial value for the states.

Parameters¶

model : Model Model to be solved exogenous : array-like, dict, or pandas.DataFrame A specification for the path of exogenous variables (aka shocks). Can be any of the following (note by "declaration order" below we mean the order of model.symbols["exogenous"]):

- A 1d numpy array-like specifying a time series for a single
  exogenous variable, or all exogenous variables stacked into a single array.
- A 2d numpy array where each column specifies the time series
  for one of the shocks in declaration order. This must be an
  `N` by number of shocks 2d array.
- A dict where keys are strings found in
  `model.symbols["exogenous"]` and values are a time series of
  values for that shock. For exogenous variables that do not appear in
  this dict, the shock is set to the calibrated value. Note
  that this interface is the most flexible as it allows the user
  to pass values for only a subset of the model shocks and it
  allows the passed time series to be of different lengths.
- A DataFrame where columns map shock names into time series.
  The same assumptions and behavior that are used in the dict
  case apply here

If nothing is given here, `exogenous` is set equal to the
calibrated values found in `model.calibration["exogenous"]` for
all periods.

If the length of any time-series in shocks is less than `T`
(see below) it is assumed that that particular shock will
remain at the final given value for the duration of the
simulation.

s0 : None or ndarray or dict If vector with the value of initial states If an exogenous timeseries is given for exogenous shocks, s0 will be computed as the steady-state value that is consistent with its first value.

int

horizon for the perfect foresight simulation

maxit : int maximum number of iteration for the nonlinear solver verbose : boolean if True, the solver displays iterations tol : float stopping criterium for the nonlinear solver ignore_constraints : bool if True, complementarity constraints are ignored. keep_steady_state : bool if True, initial steady-states and steady-controls are appended to the simulation with date -1. Returns

pandas dataframe a dataframe with T+1 observations of the model variables along the simulation (states, controls, auxiliaries). The simulation should return to a steady-state consistent with the last specified value of the exogenous shocks.

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

def deterministic_solve(
    model: Model,
    *,  #
    verbose: bool = True,  #
    ignore_constraints: bool = False,  #
    exogenous=None,
    s0=None,
    m0=None,
    T=100,
    maxit=100,
    initial_guess=None,
    solver="ncpsolve",
    keep_steady_state=False,
    s1=None,  # deprecated
    shocks=None,  # deprecated
    tol=1e-6,
):
    """
    Computes a perfect foresight simulation using a stacked-time algorithm.

    Typical simulation exercises are:
    - start from an out-of-equilibrium exogenous and/or endogenous state: specify `s0` and or `m0`. Missing values are taken from the calibration (`model.calibration`).
    - specify an exogenous path for shocks `exogenous`. Initial exogenous state `m0` is then first value of exogenous values. Economy is supposed to have been at the equilibrium for $t<0$, which pins
    down initial endogenous state `s0`. `x0` is a jump variable.

    If $s0$ is not specified it is then set
    equal to the steady-state consistent with the first value

    The initial state is specified either by providing a series of exogenous
    shocks and assuming the model is initially in equilibrium with the first
    value of the shock, or by specifying an initial value for the states.

    Parameters
    ----------
    model : Model
        Model to be solved
    exogenous : array-like, dict, or pandas.DataFrame
        A specification for the path of exogenous variables (aka shocks). Can be any of the
        following (note by "declaration order" below we mean the order
        of `model.symbols["exogenous"]`):

        - A 1d numpy array-like specifying a time series for a single
          exogenous variable, or all exogenous variables stacked into a single array.
        - A 2d numpy array where each column specifies the time series
          for one of the shocks in declaration order. This must be an
          `N` by number of shocks 2d array.
        - A dict where keys are strings found in
          `model.symbols["exogenous"]` and values are a time series of
          values for that shock. For exogenous variables that do not appear in
          this dict, the shock is set to the calibrated value. Note
          that this interface is the most flexible as it allows the user
          to pass values for only a subset of the model shocks and it
          allows the passed time series to be of different lengths.
        - A DataFrame where columns map shock names into time series.
          The same assumptions and behavior that are used in the dict
          case apply here

        If nothing is given here, `exogenous` is set equal to the
        calibrated values found in `model.calibration["exogenous"]` for
        all periods.

        If the length of any time-series in shocks is less than `T`
        (see below) it is assumed that that particular shock will
        remain at the final given value for the duration of the
        simulation.
    s0 : None or ndarray or dict
        If vector with the value of initial states
        If an exogenous timeseries is given for exogenous shocks, `s0` will be computed as the steady-state value that is consistent with its first value.

    T : int
        horizon for the perfect foresight simulation
    maxit : int
        maximum number of iteration for the nonlinear solver
    verbose : boolean
        if True, the solver displays iterations
    tol : float
        stopping criterium for the nonlinear solver
    ignore_constraints : bool
        if True, complementarity constraints are ignored.
    keep_steady_state : bool
        if True, initial steady-states and steady-controls are appended to the simulation with date -1.
    Returns
    -------
    pandas dataframe
        a dataframe with T+1 observations of the model variables along the
        simulation (states, controls, auxiliaries). The  simulation should return to a steady-state
        consistent with the last specified value of the exogenous shocks.

    """

    if shocks is not None:
        import warnings

        warnings.warn("`shocks` argument is deprecated. Use `exogenous` instead.")
        exogenous = shocks

    if s1 is not None:
        import warnings

        warnings.warn("`s1` argument is deprecated. Use `s0` instead.")
        s0 = s1

    # definitions
    n_s = len(model.calibration["states"])
    n_x = len(model.calibration["controls"])
    p = model.calibration["parameters"]

    if exogenous is not None:
        epsilons = _shocks_to_epsilons(model, exogenous, T)
        m0 = epsilons[0, :]
        # get initial steady-state
        from dolo.algos.steady_state import find_steady_state

        start_state = find_steady_state(model, m=m0)
        s0 = start_state["states"]
        x0 = start_state["controls"]
        m1 = epsilons[1, :]
        s1 = model.functions["transition"](m0, s0, x0, m1, p)
    else:
        if s0 is None:
            s0 = model.calibration["states"]
        if m0 is None:
            m0 = model.calibration["exogenous"]
        # if m0 is None:
        #     m0 = np.zeros(len(model.symbols['exogenous']))
        # we should probably do something here with the nature of the exogenous process if specified
        # i.e. compute nonlinear irf
        epsilons = _shocks_to_epsilons(model, exogenous, T)
        x0 = model.calibration["controls"]
        m1 = epsilons[1, :]
        s1 = model.functions["transition"](m0, s0, x0, m1, p)
        s1 = np.array(s1)

    x1_g = model.calibration["controls"]  # we can do better here
    sT_g = model.calibration["states"]  # we can do better here
    xT_g = model.calibration["controls"]  # we can do better here

    if initial_guess is None:
        start = np.concatenate([s1, x1_g])
        final = np.concatenate([sT_g, xT_g])
        initial_guess = np.vstack(
            [start * (1 - l) + final * l for l in linspace(0.0, 1.0, T + 1)]
        )

    else:
        if isinstance(initial_guess, pd.DataFrame):
            initial_guess = np.array(
                initial_guess[model.symbols["states"] + model.symbols["controls"]]
            )
        initial_guess = initial_guess[1:, :]
        initial_guess = initial_guess[:, : n_s + n_x]

    sh = initial_guess.shape

    if model.x_bounds and not ignore_constraints:
        initial_states = initial_guess[:, :n_s]
        [lb, ub] = [u(epsilons[:, :], initial_states, p) for u in model.x_bounds]
        lower_bound = initial_guess * 0 - np.inf
        lower_bound[:, n_s:] = lb
        upper_bound = initial_guess * 0 + np.inf
        upper_bound[:, n_s:] = ub
        test1 = max(lb.max(axis=0) - lb.min(axis=0))
        test2 = max(ub.max(axis=0) - ub.min(axis=0))
        if test1 > 0.00000001 or test2 > 0.00000001:
            msg = "Not implemented: perfect foresight solution requires that "
            msg += "controls have constant bounds."
            raise Exception(msg)
    else:
        ignore_constraints = True
        lower_bound = None
        upper_bound = None

    det_residual(model, initial_guess, s1, xT_g, epsilons)

    if not ignore_constraints:

        def ff(vec):
            return det_residual(
                model, vec.reshape(sh), s1, xT_g, epsilons, jactype="sparse"
            )

        v0 = initial_guess.ravel()
        if solver == "ncpsolve":
            sol, nit = ncpsolve(
                ff,
                lower_bound.ravel(),
                upper_bound.ravel(),
                initial_guess.ravel(),
                verbose=verbose,
                maxit=maxit,
                tol=tol,
                jactype="sparse",
            )
        else:
            from dolo.numeric.extern.lmmcp import lmmcp

            sol = lmmcp(
                lambda u: ff(u)[0],
                lambda u: ff(u)[1].todense(),
                lower_bound.ravel(),
                upper_bound.ravel(),
                initial_guess.ravel(),
                verbose=verbose,
            )
            nit = -1

        sol = sol.reshape(sh)

    else:

        def ff(vec):
            ll = det_residual(model, vec.reshape(sh), s1, xT_g, epsilons, diff=True)
            return ll

        v0 = initial_guess.ravel()
        # from scipy.optimize import root
        # sol = root(ff, v0, jac=True)
        # sol = sol.x.reshape(sh)
        from dolo.numeric.optimize.newton import newton

        sol, nit = newton(ff, v0, jactype="sparse")
        sol = sol.reshape(sh)

    # sol = sol[:-1, :]

    if (exogenous is not None) and keep_steady_state:
        sx = np.concatenate([s0, x0])
        sol = np.concatenate([sx[None, :], sol], axis=0)
        epsilons = np.concatenate([epsilons[:1:], epsilons], axis=0)
        index = range(-1, T + 1)
    else:
        index = range(0, T + 1)
    # epsilons = np.concatenate([epsilons[:1,:], epsilons], axis=0)

    if "auxiliary" in model.functions:
        colnames = (
            model.symbols["states"]
            + model.symbols["controls"]
            + model.symbols["auxiliaries"]
        )
        # compute auxiliaries
        y = model.functions["auxiliary"](epsilons, sol[:, :n_s], sol[:, n_s:], p)
        sol = np.column_stack([sol, y])
    else:
        colnames = model.symbols["states"] + model.symbols["controls"]

    sol = np.column_stack([sol, epsilons])
    colnames = colnames + model.symbols["exogenous"]

    ts = pd.DataFrame(sol, columns=colnames, index=index)
    return ts

`evaluate_policy(model, mdr, tol=1e-08, maxit=2000, grid={}, verbose=True, dr0=None, hook=None, integration_orders=None, details=False, interp_method='cubic')` ¶

Compute value function corresponding to policy dr

Parameters:¶

model

"dtcscc" model. Must contain a 'value' function.

mdr

decision rule to evaluate

Returns:¶

decision rule

value function (a function of the space similar to a decision rule object)

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

def evaluate_policy(
    model,
    mdr,
    tol=1e-8,
    maxit=2000,
    grid={},
    verbose=True,
    dr0=None,
    hook=None,
    integration_orders=None,
    details=False,
    interp_method="cubic",
):
    """Compute value function corresponding to policy ``dr``

    Parameters:
    -----------

    model:
        "dtcscc" model. Must contain a 'value' function.

    mdr:
        decision rule to evaluate

    Returns:
    --------

    decision rule:
        value function (a function of the space similar to a decision rule
        object)

    """

    process = model.exogenous
    grid, dprocess = model.discretize()
    endo_grid = grid["endo"]
    exo_grid = grid["exo"]

    n_ms = dprocess.n_nodes  # number of exogenous states
    n_mv = dprocess.n_inodes(0)  # this assume number of integration nodes is constant

    x0 = model.calibration["controls"]
    v0 = model.calibration["values"]
    parms = model.calibration["parameters"]
    n_x = len(x0)
    n_v = len(v0)
    n_s = len(model.symbols["states"])

    if dr0 is not None:
        mdrv = dr0
    else:
        mdrv = DecisionRule(exo_grid, endo_grid, interp_method=interp_method)

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

    if isinstance(mdr, np.ndarray):
        controls = mdr
    else:
        controls = np.zeros((n_ms, N, n_x))
        for i_m in range(n_ms):
            controls[i_m, :, :] = mdr(i_m, s)

    values_0 = np.zeros((n_ms, N, n_v))
    if dr0 is None:
        for i_m in range(n_ms):
            values_0[i_m, :, :] = v0[None, :]
    else:
        for i_m in range(n_ms):
            values_0[i_m, :, :] = dr0(i_m, s)

    val = model.functions["value"]
    g = model.functions["transition"]

    sh_v = values_0.shape

    err = 10
    inner_maxit = 50
    it = 0

    if verbose:
        headline = "|{0:^4} | {1:10} | {2:8} | {3:8} |".format(
            "N", " Error", "Gain", "Time"
        )
        stars = "-" * len(headline)
        print(stars)
        print(headline)
        print(stars)

    t1 = time.time()

    err_0 = np.nan

    verbit = verbose == "full"

    while err > tol and it < maxit:

        it += 1

        t_start = time.time()

        mdrv.set_values(values_0.reshape(sh_v))
        values = update_value(
            val, g, s, controls, values_0, mdr, mdrv, dprocess, parms
        ).reshape((-1, n_v))
        err = abs(values.reshape(sh_v) - values_0).max()

        err_SA = err / err_0
        err_0 = err

        values_0 = values.reshape(sh_v)

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

        if verbose:
            print(
                "|{0:4} | {1:10.3e} | {2:8.3f} | {3:8.3f} |".format(
                    it, err, err_SA, elapsed
                )
            )

    # values_0 = values.reshape(sh_v)

    t2 = time.time()

    if verbose:
        print(stars)
        print("Elapsed: {} seconds.".format(t2 - t1))
        print(stars)

    if not details:
        return mdrv
    else:
        return EvaluationResult(mdrv, it, tol, err)

`newton(model, *, engine='auto', backend=None, device=None, dr0=None, tol_eps=1e-08, tol_eta=1e-06, maxit=30, linsolve='gmres', linesearch='backtracking', interpolation='cubic', verbose=True, details=True, trace=False, dprocess=None, ignore_constraints=False, compute_radius=False, **kwargs)` ¶

Finds a global solution for model using full-system Newton-Krylov solver.

Parameters¶

model : Model model to be solved 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 (30) maximum number of outer Newton iterations linsolve : str or dict ("gmres") linear solver for the preconditioned system (I - L) dx = r. Options: "gmres", "neumann", or a dict(method=..., maxit=..., tol=...) linesearch : str or dict ("backtracking") step size stabilization. Options: "backtracking", "golden_section", False, None, or a dict(method=..., max_backsteps=...) 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/newton.py

def newton(
    model: Model,
    *,
    engine: str = "auto",
    backend: Optional[str] = None,
    device: Optional[Any] = None,
    dr0: DecisionRule = None,
    tol_eps: float = 1e-8,
    tol_eta: float = 1e-6,
    maxit: int = 30,
    linsolve="gmres",
    linesearch="backtracking",
    interpolation: str = "cubic",
    verbose: bool = True,
    details: bool = True,
    trace: bool = False,
    dprocess=None,
    ignore_constraints: bool = False,
    compute_radius: bool = False,
    **kwargs,
) -> NewtonResult:
    """Finds a global solution for ``model`` using full-system Newton-Krylov solver.

    Parameters
    ----------
    model : Model
        model to be solved
    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 (30)
        maximum number of outer Newton iterations
    linsolve : str or dict ("gmres")
        linear solver for the preconditioned system (I - L) dx = r.
        Options: "gmres", "neumann", or a dict(method=..., maxit=..., tol=...)
    linesearch : str or dict ("backtracking")
        step size stabilization. Options: "backtracking", "golden_section", False, None,
        or a dict(method=..., max_backsteps=...)
    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_eps = kwargs.pop("tol")
    if "interp_method" in kwargs:
        interpolation = kwargs.pop("interp_method")

    complementarities = kwargs.pop("complementarities", None)
    invmethod = kwargs.pop("invmethod", None)
    smaxit = kwargs.pop("smaxit", None)
    linear_maxit = kwargs.pop("linear_maxit", None)
    max_backsteps = kwargs.pop("max_backsteps", None)
    maxbsteps = kwargs.pop("maxbsteps", None)
    mu = kwargs.pop("mu", 2)
    method = kwargs.pop("method", "jac")

    if complementarities is not None:
        pass
    else:
        complementarities = not ignore_constraints

    # Linear solver config
    if isinstance(linsolve, str):
        linsolve_cfg = {
            "method": "neumann" if linsolve == "iti" else linsolve,
            "maxit": 50,
            "tol": 1e-8,
            "precondition": True,
        }
    elif isinstance(linsolve, dict):
        linsolve_cfg = {
            "method": "gmres",
            "maxit": 50,
            "tol": 1e-8,
            "precondition": True,
            **linsolve,
        }
        if linsolve_cfg["method"] == "iti":
            linsolve_cfg["method"] = "neumann"
    else:
        linsolve_cfg = {
            "method": "gmres",
            "maxit": 50,
            "tol": 1e-8,
            "precondition": True,
        }

    if invmethod is not None:
        linsolve_cfg["method"] = "neumann" if invmethod == "iti" else invmethod
    if smaxit is not None:
        linsolve_cfg["maxit"] = smaxit
    if linear_maxit is not None:
        linsolve_cfg["maxit"] = linear_maxit

    # Line search config
    if linesearch is False or linesearch is None or linesearch == "none":
        linesearch_cfg = {"method": None, "max_backsteps": 0}
    elif isinstance(linesearch, str):
        linesearch_cfg = {"method": linesearch, "max_backsteps": 10}
    elif isinstance(linesearch, dict):
        linesearch_cfg = {"method": "backtracking", "max_backsteps": 10, **linesearch}
    else:
        linesearch_cfg = {"method": "backtracking", "max_backsteps": 10}

    if max_backsteps is not None:
        linesearch_cfg["max_backsteps"] = max_backsteps
    elif maxbsteps is not None:
        linesearch_cfg["max_backsteps"] = maxbsteps

    itprint = IterationsPrinter(
        ("N", int),
        ("f_x", float),
        ("d_x", float),
        ("Time_residuals", float),
        ("Time_inversion", float),
        ("Time_search", float),
        ("Lambda_0", float),
        ("N_invert", int),
        ("N_search", int),
        verbose=verbose,
    )
    itprint.print_header("Start Newton Iterations.")

    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

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

    f = dmodel.f
    g = dmodel.g
    x_lb = dmodel.model.functions.get("arbitrage_lb", None)
    x_ub = dmodel.model.functions.get("arbitrage_ub", None)

    parms = dmodel.parms

    n_m = max(dprocess.n_nodes, 1)
    n_s = len(dmodel.model.symbols["states"])

    eng = _resolve_engine(dmodel, engine, backend)
    use_jax = eng in ("jax", "gpu", "cpu") and n_m == 1
    if use_jax:
        target_dev = device
        if target_dev is None:
            if eng == "gpu":
                import jax

                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":
                import jax

                target_dev = jax.devices("cpu")[0]
    else:
        target_dev = None

    if interpolation in ("cubic", "linear"):
        ddr = DecisionRule(
            exo_grid, endo_grid, dprocess=dprocess, interp_method=interpolation
        )
        ddr_filt = DecisionRule(
            exo_grid, endo_grid, dprocess=dprocess, interp_method=interpolation
        )
    else:
        raise Exception("Unsupported interpolation method.")

    s = endo_grid.nodes
    N = s.shape[0]
    n_x = len(dmodel.model.symbols["controls"])
    x0 = (
        dmodel.model.calibration["controls"][
            None,
            None,
        ]
        .repeat(n_m, axis=0)
        .repeat(N, axis=1)
    )

    if dr0 is not None:
        for i_m in range(n_m):
            x0[i_m, :, :] = dr0(i_m, s)
    ddr.set_values(x0)

    if x_lb is not None and complementarities:
        lb = x0.copy()
        ub = x0.copy()
        for i_m in range(n_m):
            m = dprocess.node(i_m)
            lb[i_m, :] = x_lb(m, s, parms)
            ub[i_m, :] = x_ub(m, s, parms)
    else:
        complementarities = False
        lb = None
        ub = None

    x = x0.copy()
    ddr.set_values(x)

    n_im = dprocess.n_inodes(0)
    jres = numpy.zeros((n_m, n_im, N, n_x, n_x))
    S_ij = numpy.zeros((n_m, n_im, N, n_s))

    trace_details = [] if trace else None

    err_0 = 10.0
    err_2 = 10.0

    for it in range(maxit):
        if use_jax and not complementarities:
            import jax
            import jax.numpy as jnp

            t1 = time.time()
            ddr.set_values(x)

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

            # 1. Residuals and Jacobian
            res, dres = F_and_dF_1(
                dmodel,
                x,
                dr=ddr,
                engine=eng,
                interp_method=interpolation,
                device=target_dev,
            )
            res_arr = numpy.asarray(
                res.data[0]
                if isinstance(res, GArray)
                else (res[0] if res.ndim == 3 else res)
            )
            dres_arr = numpy.asarray(
                dres.data[0]
                if isinstance(dres, GArray)
                else (dres[0] if dres.ndim == 4 else dres)
            )

            err_0 = abs(res_arr).max()
            if err_0 < tol_eps:
                t2 = time.time()
                t3 = t2
                t4 = t2
                itprint.print_iteration(
                    N=it,
                    f_x=err_0,
                    d_x=0.0,
                    Time_residuals=t2 - t1,
                    Time_inversion=0.0,
                    Time_search=0.0,
                    Lambda_0=0.0,
                    N_invert=0,
                    N_search=0,
                )
                break

            t2 = time.time()

            # 2. Block-diagonal solve: ddx = dres^{-1} * res
            J = jnp.asarray(dres_arr)
            R = jnp.asarray(res_arr)
            ddx = numpy.asarray(jnp.linalg.solve(J, R[..., None])[..., 0])

            # 3. Policy linear operator dF_2 and preconditioned inversion
            solver_method = linsolve_cfg.get("method", "gmres")
            solver_maxit = linsolve_cfg.get("maxit", 50)
            solver_tol = linsolve_cfg.get("tol", 1e-8)
            ttol = max(err_0 / 100, solver_tol)

            L_op = dF_2(
                dmodel,
                x,
                dr=ddr,
                engine=eng,
                interp_method=interpolation,
                device=target_dev,
            )
            L_prec = (-1.0) * L_op.preconditioned(J)
            tot_2d = L_prec.solve(
                ddx, linsolve=solver_method, tol=ttol, maxit=solver_maxit
            )
            tot = numpy.asarray(tot_2d).reshape(x.shape)
            L_last = L_prec
            nn = solver_maxit
            lam0 = 0.01

            t3 = time.time()

            # 4. Line search
            ls_method = linesearch_cfg.get("method", "backtracking")
            ls_maxsteps = linesearch_cfg.get("max_backsteps", 10)

            def eval_res_norm_jax(cand_x):
                c_res = F(
                    dmodel,
                    cand_x,
                    dr=ddr,
                    engine=eng,
                    interp_method=interpolation,
                    device=target_dev,
                )
                c_arr = numpy.asarray(
                    c_res.data[0]
                    if isinstance(c_res, GArray)
                    else (c_res[0] if c_res.ndim == 3 else c_res)
                )
                return abs(c_arr).max()

            if ls_method == "golden_section":
                cand_1 = eval_res_norm_jax(x - tot)
                if cand_1 < err_0:
                    best_alpha = 1.0
                    i_bckstps = 0
                else:
                    best_alpha, _ = golden_section_search(
                        lambda a: eval_res_norm_jax(x - tot * a),
                        alpha_min=0.0,
                        alpha_max=1.0,
                        maxiter=ls_maxsteps,
                    )
                    i_bckstps = ls_maxsteps
                new_x = x - tot * best_alpha

            elif ls_method == "backtracking":
                steps = 0.5 ** numpy.arange(max(ls_maxsteps, 1))
                new_x = x - tot
                i_bckstps = 0
                for i_bckstps, lam in enumerate(steps):
                    new_x = x - tot * lam
                    new_err = eval_res_norm_jax(new_x)
                    if new_err < err_0:
                        break
            else:
                new_x = x - tot
                i_bckstps = 0

            err_2 = abs(tot).max()
            t4 = time.time()

            itprint.print_iteration(
                N=it,
                f_x=err_0,
                d_x=err_2,
                Time_residuals=t2 - t1,
                Time_inversion=t3 - t2,
                Time_search=t4 - t3,
                Lambda_0=lam0,
                N_invert=nn,
                N_search=i_bckstps,
            )

            x = new_x
            if err_0 < tol_eps or err_2 < tol_eta:
                break

        else:
            jres[...] = 0.0
            S_ij[...] = 0.0

            t1 = time.time()
            ddr.set_values(x)

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

            sh_x = x.shape
            ff = SerialDifferentiableFunction(
                lambda u: euler_residuals(
                    f,
                    g,
                    s,
                    u.reshape(sh_x),
                    ddr,
                    dprocess,
                    parms,
                    diff=False,
                    with_jres=False,
                    set_dr=False,
                ).reshape((-1, sh_x[2]))
            )
            res, dres = ff(x.reshape((-1, sh_x[2])))
            res = res.reshape(sh_x)
            dres = dres.reshape((sh_x[0], sh_x[1], sh_x[2], sh_x[2]))
            junk, jres, fut_S = euler_residuals(
                f,
                g,
                s,
                x,
                ddr,
                dprocess,
                parms,
                diff=False,
                with_jres=True,
                set_dr=False,
                jres=jres,
                S_ij=S_ij,
            )

            if complementarities:
                res, dres, jres = smooth(res, dres, jres, x - lb)
                res[...] *= -1
                dres[...] *= -1
                jres[...] *= -1
                res, dres, jres = smooth(res, dres, jres, ub - x, pos=-1.0)
                res[...] *= -1
                dres[...] *= -1
                jres[...] *= -1

            err_0 = abs(res).max()

            jres[...] *= -1.0

            for i_m in range(n_m):
                for j_m in range(n_im):
                    M = jres[i_m, j_m, :, :, :]
                    X = dres[i_m, :, :, :].copy()
                    sol = solve_tensor(X, M)

            t2 = time.time()

            # Inversion of (I - L) dx = F_1^{-1} res
            solver_method = linsolve_cfg.get("method", "gmres")
            solver_maxit = linsolve_cfg.get("maxit", 50)
            solver_tol = linsolve_cfg.get("tol", 1e-8)

            if solver_method == "gmres":
                ddx = solve_gu(dres.copy(), res.copy())
                L = Operator(jres, fut_S, ddr_filt)
                n0 = L.counter
                L.addid = True
                ttol = max(err_0 / 100, solver_tol)
                try:
                    sol, info = scipy.sparse.linalg.gmres(
                        L, ddx.ravel(), rtol=ttol, atol=0.0, maxiter=solver_maxit
                    )
                except TypeError:
                    sol, info = scipy.sparse.linalg.gmres(
                        L, ddx.ravel(), tol=ttol, maxiter=solver_maxit
                    )
                lam0 = 0.01
                nn = L.counter - n0
                tot = sol.reshape(ddx.shape)
            else:
                # Neumann series
                tot, nn, lam0 = invert_jac(
                    res,
                    dres,
                    jres,
                    fut_S,
                    ddr_filt,
                    tol=solver_tol,
                    maxit=solver_maxit,
                    verbose=(verbose == "full"),
                )

            t3 = time.time()

            # Step search / Line search
            ls_method = linesearch_cfg.get("method", "backtracking")
            ls_maxsteps = linesearch_cfg.get("max_backsteps", 10)

            def eval_residual_norm(cand_x):
                cand_res = euler_residuals(
                    f, g, s, cand_x, ddr, dprocess, parms, diff=False, set_dr=True
                )
                if complementarities:
                    cand_res = smooth_nodiff(cand_res, cand_x - lb)
                    cand_res = smooth_nodiff(-cand_res, ub - cand_x)
                return abs(cand_res).max()

            if ls_method == "golden_section":
                # Check full step first
                cand_1 = eval_residual_norm(x - tot)
                if cand_1 < err_0:
                    best_alpha = 1.0
                    i_bckstps = 0
                else:
                    best_alpha, _ = golden_section_search(
                        lambda a: eval_residual_norm(x - tot * a),
                        alpha_min=0.0,
                        alpha_max=1.0,
                        maxiter=ls_maxsteps,
                    )
                    i_bckstps = ls_maxsteps
                new_x = x - tot * best_alpha

            elif ls_method == "backtracking":
                steps = 0.5 ** numpy.arange(max(ls_maxsteps, 1))
                new_x = x - tot
                for i_bckstps, lam in enumerate(steps):
                    new_x = x - tot * lam
                    new_err = eval_residual_norm(new_x)
                    if new_err < err_0:
                        break
            else:
                # No line search (full Newton step)
                new_x = x - tot
                i_bckstps = 0

            err_2 = abs(tot).max()
            t4 = time.time()

            itprint.print_iteration(
                N=it,
                f_x=err_0,
                d_x=err_2,
                Time_residuals=t2 - t1,
                Time_inversion=t3 - t2,
                Time_search=t4 - t3,
                Lambda_0=lam0,
                N_invert=nn,
                N_search=i_bckstps,
            )

            if err_0 < tol_eps:
                break

            x = new_x

    ddr.set_values(x)
    if trace:
        trace_details.append({"dr": copy.deepcopy(ddr)})

    itprint.print_finished()

    if not details:
        return ddr
    else:
        if use_jax and not complementarities:
            L = L_last if "L_last" in locals() else None
            if compute_radius and L is not None:
                try:
                    lam_eig = scipy.sparse.linalg.eigs(L, k=1, return_eigenvectors=False)
                    lam = abs(lam_eig[0])
                except Exception:
                    lam = numpy.nan
            else:
                lam = numpy.nan
        else:
            L = Operator(jres, fut_S, ddr_filt)
            if compute_radius:
                try:
                    lam_eig = scipy.sparse.linalg.eigs(L, k=1, return_eigenvectors=False)
                    lam = abs(lam_eig[0])
                except Exception:
                    lam = numpy.nan
            else:
                lam = numpy.nan

        return NewtonResult(
            ddr,
            it + 1,
            err_0,
            err_2,
            err_0 < tol_eps,
            complementarities,
            lam,
            trace_details,
            L,
        )

`perturb(model, *, details=True, verbose=True, steady_state=None, eigmax=1.0 - 1e-06, solve_steady_state=False, order=1)` ¶

Compute first order approximation of optimal controls

Parameters:¶

NumericModel

Model to be solved

boolean

If True: displays number of contracting eigenvalues

ndarray

Use supplied steady-state value to compute the approximation. The routine doesn't check whether it is really a solution or not.

boolean

Use nonlinear solver to find the steady-state

{1}

Approximation order. (Currently, only first order is supported).

Returns:¶

TaylorExpansion

Decision Rule for the optimal controls around the steady-state.

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

def perturb(
    model: Model,
    *,
    details: bool = True,  #
    verbose: bool = True,  #
    steady_state=None,
    eigmax=1.0 - 1e-6,
    solve_steady_state=False,
    order=1,
):
    """Compute first order approximation of optimal controls

    Parameters:
    -----------

    model: NumericModel
        Model to be solved

    verbose: boolean
        If True: displays number of contracting eigenvalues

    steady_state: ndarray
        Use supplied steady-state value to compute the approximation.
        The routine doesn't check whether it is really a solution or not.

    solve_steady_state: boolean
        Use nonlinear solver to find the steady-state

    orders: {1}
        Approximation order. (Currently, only first order is supported).

    Returns:
    --------

    TaylorExpansion:
        Decision Rule for the optimal controls around the steady-state.

    """

    if order > 1:
        raise Exception("Not implemented.")

    if steady_state is None:
        steady_state = model.calibration

    G_s, G_x, G_e, F_s, F_x, F_S, F_X = get_derivatives(
        model, steady_state=steady_state
    )

    C, eigvals = approximate_1st_order(G_s, G_x, G_e, F_s, F_x, F_S, F_X)

    m = steady_state["exogenous"]
    s = steady_state["states"]
    x = steady_state["controls"]

    from dolo.numeric.processes import VAR1, IIDProcess
    from dolo.numeric.distribution import MvNormal

    process = model.exogenous

    if isinstance(process, VAR1):
        C_m = C[:, : len(m)]
        C_s = C[:, len(m) :]
    elif isinstance(process, IIDProcess):
        C_m = None
        C_s = C
    dr = BivariateTaylor(m, s, x, C_m, C_s)
    if not details:
        return dr
    else:
        return PerturbationResult(
            dr,
            eigvals,
            True,  # otherwise an Exception should have been raised already
            True,  # otherwise an Exception should have been raised already
            True,  # otherwise an Exception should have been raised already
        )

`simulate(model, dr, *, process=None, N=1, T=40, s0=None, i0=None, m0=None, driving_process=None, seed=42, stochastic=True)` ¶

Simulate a model using the specified decision rule.

Parameters¶

model: Model

dr: decision rule

process:

ndarray

initial state where all simulations start

ndarray

realization of exogenous driving process (drawn randomly if None)

int

number of simulations

T: int horizon for the simulations seed: int used to initialize the random number generator. Use it to replicate exact same results among simulations discard: boolean (False) if True, then all simulations containing at least one non finite value are discarded

Returns¶

xarray.DataArray: returns a T x N x n_v array where n_v is the number of variables.

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

def simulate(
    model: Model,
    dr: DecisionRule,
    *,
    process=None,
    N=1,
    T=40,
    s0=None,
    i0=None,
    m0=None,
    driving_process=None,
    seed=42,
    stochastic=True,
):
    """Simulate a model using the specified decision rule.

    Parameters
    ----------

    model: Model

    dr: decision rule

    process:

    s0: ndarray
        initial state where all simulations start

    driving_process: ndarray
        realization of exogenous driving process (drawn randomly if None)

    N: int
        number of simulations
    T: int
        horizon for the simulations
    seed: int
        used to initialize the random number generator. Use it to replicate
        exact same results among simulations
    discard: boolean (False)
        if True, then all simulations containing at least one non finite value
        are discarded

    Returns
    -------
    xarray.DataArray:
        returns a ``T x N x n_v`` array where ``n_v``
           is the number of variables.
    """

    if isinstance(dr, AlgoResult):
        dr = dr.dr

    calib = model.calibration
    parms = numpy.array(calib["parameters"])

    if s0 is None:
        s0 = calib["states"]

    n_x = len(model.symbols["controls"])
    n_s = len(model.symbols["states"])

    s_simul = numpy.zeros((T, N, n_s))
    x_simul = numpy.zeros((T, N, n_x))

    s_simul[0, :, :] = s0[None, :]

    # are we simulating a markov chain or a continuous process ?
    if driving_process is not None:
        if len(driving_process.shape) == 3:
            m_simul = driving_process
            sim_type = "continuous"
            if m0 is None:
                m0 = model.calibration["exogenous"]
            x_simul[0, :, :] = dr(m0[None, :], s0[None, :])[0, :]
        elif len(driving_process.shape) == 2:
            i_simul = driving_process
            nodes = dr.exo_grid.nodes
            m_simul = nodes[i_simul]
            # inds = i_simul.ravel()
            # m_simul = np.reshape( np.concatenate( [nodes[i,:][None,:] for i in inds.ravel()], axis=0 ), inds.shape + (-1,) )
            sim_type = "discrete"
            x_simul[0, :, :] = dr(i0, s0[None, :])[0, :]
        else:
            raise Exception("Incorrect specification of driving values.")
        m0 = m_simul[0, :, :]
    else:
        from dolo.numeric.processes import DiscreteProcess

        if process is None:
            if hasattr(dr, "dprocess") and hasattr(dr.dprocess, "simulate"):
                process = dr.dprocess
            else:
                process = model.exogenous

        # detect type of simulation
        if not isinstance(process, DiscreteProcess):
            sim_type = "continuous"
        else:
            sim_type = "discrete"

        if sim_type == "discrete":
            if i0 is None:
                i0 = 0
            dp = process
            m_simul = dp.simulate(N, T, i0=i0, stochastic=stochastic)
            i_simul = find_index(m_simul, dp.values)
            m0 = dp.node(i0)
            x0 = dr(i0, s0[None, :])[0, :]
        else:
            m_simul = process.simulate(N, T, m0=m0, stochastic=stochastic)
            if isinstance(m_simul, xr.DataArray):
                m_simul = m_simul.data
            sim_type = "continuous"
            if m0 is None:
                m0 = model.calibration["exogenous"]
            x0 = dr(m0[None, :], s0[None, :])[0, :]
            x_simul[0, :, :] = x0[None, :]

    f = model.functions["arbitrage"]
    g = model.functions["transition"]

    numpy.random.seed(seed)

    mp = m0
    for i in range(T):
        m = m_simul[i, :, :]
        s = s_simul[i, :, :]
        if sim_type == "discrete":
            i_m = i_simul[i, :]
            xx = [dr(i_m[ii], s[ii, :][None, :])[0, :] for ii in range(s.shape[0])]
            x = np.vstack(xx)
        else:
            x = dr(m, s)

        x_simul[i, :, :] = x

        ss = g(mp, s, x, m, parms)
        if i < T - 1:
            s_simul[i + 1, :, :] = ss
        mp = m

    if "auxiliary" not in model.functions:  # TODO: find a better test than this
        l = [s_simul, x_simul]
        varnames = model.symbols["states"] + model.symbols["controls"]
    else:
        aux = model.functions["auxiliary"]
        a_simul = aux(
            m_simul.reshape((N * T, -1)),
            s_simul.reshape((N * T, -1)),
            x_simul.reshape((N * T, -1)),
            parms,
        )
        a_simul = a_simul.reshape(T, N, -1)
        l = [m_simul, s_simul, x_simul, a_simul]
        varnames = (
            model.symbols["exogenous"]
            + model.symbols["states"]
            + model.symbols["controls"]
            + model.symbols["auxiliaries"]
        )

    simul = numpy.concatenate(l, axis=2)

    if sim_type == "discrete":
        varnames = ["_i_m"] + varnames
        simul = np.concatenate([i_simul[:, :, None], simul], axis=2)

    data = xr.DataArray(
        simul,
        dims=["T", "N", "V"],
        coords={"T": range(T), "N": range(N), "V": varnames},
    )

    return data

`time_iteration(model, *, engine='auto', backend=None, dr0=None, tol_eps=1e-08, tol_eta=1e-06, maxit=500, inner=None, accelerate=False, interpolation='cubic', verbose=True, details=True, trace=False, ignore_constraints=False, hook=None, dprocess=None, **kwargs)` ¶

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,
    )

`value_iteration(model, *, verbose=False, details=True, tol=1e-06, maxit=500, maxit_howard=20)` ¶

Solve for the value function and associated Markov decision rule by iterating over the value function.

Parameters:¶

model : model to be solved dr : decision rule to evaluate

Returns:¶

mdr : Markov decision rule The solved decision rule/policy function mdrv: decision rule The solved value function

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

def value_iteration(
    model: Model,
    *,
    verbose: bool = False,  #
    details: bool = True,  #
    tol=1e-6,
    maxit=500,
    maxit_howard=20,
) -> ValueIterationResult:
    """
    Solve for the value function and associated Markov decision rule by iterating over
    the value function.

    Parameters:
    -----------
    model :
        model to be solved
    dr :
        decision rule to evaluate

    Returns:
    --------
    mdr : Markov decision rule
        The solved decision rule/policy function
    mdrv: decision rule
        The solved value function
    """

    transition = model.functions["transition"]
    felicity = model.functions["felicity"]
    # bounds come from the model adapters: explicit `controls_lb`/`controls_ub`
    # sections, or the bounds extracted from arbitrage complementarities.
    x_bounds = model.x_bounds
    if x_bounds is None:
        raise Exception(
            "value_iteration requires control bounds: declare "
            "`controls_lb`/`controls_ub` sections or arbitrage "
            "complementarities."
        )
    controls_lb, controls_ub = x_bounds

    parms = model.calibration["parameters"]
    discount = model.calibration["beta"]

    x0 = model.calibration["controls"]
    m0 = model.calibration["exogenous"]
    s0 = model.calibration["states"]
    r0 = felicity(m0, s0, x0, parms)

    process = model.exogenous

    grid, dprocess = model.discretize()
    endo_grid = grid["endo"]
    exo_grid = grid["exo"]

    n_ms = dprocess.n_nodes  # number of exogenous states
    n_mv = dprocess.n_inodes(0)  # this assume number of integration nodes is constant

    mdrv = DecisionRule(exo_grid, endo_grid)

    s = mdrv.endo_grid.nodes
    N = s.shape[0]
    n_x = len(x0)

    mdr = constant_policy(model)

    controls_0 = np.zeros((n_ms, N, n_x))
    p_batched = np.repeat(parms[None, :], N, axis=0)
    for i_ms in range(n_ms):
        m = np.repeat(dprocess.node(i_ms)[None, :], N, axis=0)
        lb_i = controls_lb(m, s, p_batched)
        ub_i = controls_ub(m, s, p_batched)
        x_i = mdr(i_ms, s)
        # the calibrated controls may violate the bounds: project onto
        # [lb, ub] (feasible constant policy, keeps transitioned states sane)
        controls_0[i_ms, :, :] = np.clip(x_i, lb_i, ub_i)

    values_0 = np.zeros((n_ms, N, 1))
    # for i_ms in range(n_ms):
    #     values_0[i_ms, :, :] = mdrv(i_ms, grid)

    mdr = DecisionRule(exo_grid, endo_grid)
    # mdr.set_values(controls_0)

    # THIRD: value function iterations until convergence
    it = 0
    err_v = 100
    err_v_0 = 0
    gain_v = 0.0
    err_x = 100
    err_x_0 = 0
    tol_x = 1e-5
    tol_v = 1e-7

    itprint = IterationsPrinter(
        ("N", int),
        ("Error_V", float),
        ("Gain_V", float),
        ("Error_x", float),
        ("Gain_x", float),
        ("Eval_n", int),
        ("Time", float),
        verbose=verbose,
    )
    itprint.print_header("Start value function iterations.")

    while (it < maxit) and (err_v > tol or err_x > tol_x):

        t_start = time.time()
        it += 1

        mdr.set_values(controls_0)
        if it > 2:
            ev = evaluate_policy(model, mdr, dr0=mdrv, verbose=False, details=True)
        else:
            ev = evaluate_policy(model, mdr, verbose=False, details=True)

        mdrv = ev.solution
        for i_ms in range(n_ms):
            values_0[i_ms, :, :] = mdrv(i_ms, s)

        values = values_0.copy()
        controls = controls_0.copy()

        for i_m in range(n_ms):
            m = dprocess.node(i_m)
            for n in range(N):
                s_ = s[n, :]
                x = controls[i_m, n, :]
                lb = controls_lb(m, s_, parms)
                ub = controls_ub(m, s_, parms)
                bnds = [e for e in zip(lb, ub)]

                def valfun(xx):
                    return -choice_value(
                        transition,
                        felicity,
                        i_m,
                        s_,
                        xx,
                        mdrv,
                        dprocess,
                        parms,
                        discount,
                    )[0]

                res = scipy.optimize.minimize(valfun, x, bounds=bnds)
                controls[i_m, n, :] = res.x
                values[i_m, n, 0] = -valfun(x)

        # compute error, update value and dr
        err_x = abs(controls - controls_0).max()
        err_v = abs(values - values_0).max()
        t_end = time.time()
        elapsed = t_end - t_start

        values_0 = values
        controls_0 = controls

        gain_x = err_x / err_x_0
        gain_v = err_v / err_v_0

        err_x_0 = err_x
        err_v_0 = err_v

        itprint.print_iteration(
            N=it,
            Error_V=err_v,
            Gain_V=gain_v,
            Error_x=err_x,
            Gain_x=gain_x,
            Eval_n=ev.iterations,
            Time=elapsed,
        )

    itprint.print_finished()

    mdr = DecisionRule(exo_grid, endo_grid)

    mdr.set_values(controls)
    mdrv.set_values(values_0)

    if not details:
        return mdr, mdrv
    else:
        return ValueIterationResult(
            mdr,  #:AbstractDecisionRule
            mdrv,  #:AbstractDecisionRule
            it,  #:Int
            dprocess,  #:AbstractDiscretizedProcess
            err_x < tol_x,  #:Bool
            tol_x,  #:Float64
            err_x,  #:Float64
            err_v < tol_v,  #:Bool
            tol_v,  #:Float64
            err_v,  #:Float64
            None,  # log:     #:ValueIterationLog
            None,  # trace:   #:Union{Nothing,IterationTrace
        )

Model¶

Model objects, produced by yaml_import, are the central piece of the package: they hold the parsed specification, the calibration, and the compiled functions of the model.

`Model` ¶

Bases: SymbolicModel

Model Object

Source code in src/dolo/compiler/model.py

class Model(SymbolicModel):
    """Model Object"""

    def __init__(self, data, check=True, compat=True):

        self.__compat__ = True

        super().__init__(data)

        self.model_type = "dtcc"
        self.__functions__ = None
        # self.__compile_functions__()
        self.set_changed(all="True")

        if check:
            self.symbols
            self.definitions
            self.calibration
            self.domain
            self.exogenous
            self.x_bounds
            self.functions

    def set_changed(self, all=False):
        self.__domain__ = None
        self.__exogenous__ = None
        self.__calibration__ = None
        if all:
            self.__symbols__ = None
            self.__definitions__ = None
            self.__variables__ = None
            self.__equations__ = None

    def recalibrate(self, *pargs, **kwargs):
        """Return a new, independent model with the given calibration
        overrides.

        Models are immutable (dyno semantics): the calibration of ``self``
        is left untouched.  Assign the returned model instead of mutating
        in place::

            model = model.recalibrate(delta=0.5)
        """
        if len(pargs) == 1:
            kwargs = dict(pargs[0], **kwargs)
        fresh = copy.copy(self)
        fresh.data = copy.deepcopy(self.data)
        fresh.data["calibration"].update(kwargs)
        fresh.set_changed()
        return fresh

    def set_calibration(self, *pargs, **kwargs):
        raise NotImplementedError(
            "Models are immutable: `set_calibration` has been removed.\n"
            "Recalibration returns a new model — assign it instead:\n\n"
            "    model = model.recalibrate(...)\n\n"
            "e.g. `model = model.recalibrate(delta=0.5)`."
        )

    @property
    def calibration(self):
        if self.__calibration__ is None:
            calibration_dict = super().get_calibration()
            from dolo.compiler.misc import CalibrationDict, calibration_to_vector

            calib = calibration_to_vector(self.symbols, calibration_dict)
            self.__calibration__ = CalibrationDict(self.symbols, calib)  #
        return self.__calibration__

    @property
    def exogenous(self):
        if self.__exogenous__ is None:
            self.__exogenous__ = super(self.__class__, self).get_exogenous()
        return self.__exogenous__

    @property
    def domain(self):
        if self.__domain__ is None:
            self.__domain__ = super().get_domain()
        return self.__domain__

    def discretize(self, grid_options=None, dprocess_options={}):

        dprocess = self.exogenous.discretize(**dprocess_options)

        if grid_options is None:
            endo_grid = self.endo_grid
        else:
            endo_grid = self.domain.discretize(**grid_options)

        from dolo.numeric.grids import ProductGrid

        grid = ProductGrid(dprocess.grid, endo_grid, names=["exo", "endo"])
        return [grid, dprocess]

    def __compile_functions__(self):

        from dolang.function_compiler import make_method_from_factory

        from dolang.vectorize import standard_function
        from dolo.compiler.factories import get_factory
        from .misc import LoosyDict

        equivalent_function_names = {
            "equilibrium": "arbitrage",
            "optimality": "arbitrage",
        }
        functions = LoosyDict(equivalences=equivalent_function_names)
        original_functions = {}
        original_gufunctions = {}

        funnames = [*self.equations.keys()]
        if len(self.definitions) > 0:
            funnames = funnames + ["auxiliary"]

        import dolo.config

        debug = dolo.config.debug

        for funname in funnames:

            fff = get_factory(self, funname)
            fun, gufun = make_method_from_factory(fff, vectorize=True, debug=debug)
            n_output = len(fff.content)
            functions[funname] = standard_function(gufun, n_output)
            original_gufunctions[funname] = gufun  # basic gufun function
            original_functions[funname] = fun  # basic numba fun

        self.__original_functions__ = original_functions
        self.__original_gufunctions__ = original_gufunctions
        self.__functions__ = functions

    @property
    def functions(self):
        if self.__functions__ is None:
            self.__compile_functions__()
        return self.__functions__

    def __str__(self):

        from dolo.misc.termcolor import colored
        from numpy import zeros

        s = """
        Model:
        ------
        name: "{name}"
        type: "{type}"
        file: "{filename}\n""".format(**self.infos)

        ss = "\nEquations:\n----------\n\n"
        res = self.residuals()
        res.update({"definitions": zeros(1)})

        equations = self.equations.copy()
        definitions = self.definitions
        tmp = []
        for deftype in definitions:
            tmp.append(deftype + " = " + definitions[deftype])
        definitions = {"definitions": tmp}
        equations.update(definitions)
        # for eqgroup, eqlist in self.symbolic.equations.items():
        for eqgroup in res.keys():
            if eqgroup == "auxiliary":
                continue
            if eqgroup == "definitions":
                eqlist = equations[eqgroup]
                # Update the residuals section with the right number of empty
                # values. Note: adding 'zeros' was easiest (rather than empty
                # cells), since other variable types have  arrays of zeros.
                res.update({"definitions": [None for i in range(len(eqlist))]})
            else:
                eqlist = equations[eqgroup]
            ss += "{}\n".format(eqgroup)
            for i, eq in enumerate(eqlist):
                val = res[eqgroup][i]
                if val is None:
                    ss += " {eqn:2} : {eqs}\n".format(eqn=str(i + 1), eqs=eq)
                else:
                    if abs(val) < 1e-8:
                        val = 0
                    vals = "{:.4f}".format(val)
                    if abs(val) > 1e-8:
                        vals = colored(vals, "red")
                    ss += " {eqn:2} : {vals} : {eqs}\n".format(
                        eqn=str(i + 1), vals=vals, eqs=eq
                    )
            ss += "\n"
        s += ss

        return s

    def __repr__(self):
        return self.__str__()

    def _repr_html_(self):

        from dolang.latex import eq2tex

        # general informations
        infos = self.infos
        table_infos = """
        <table>
         <td><b>Model</b></td>
        <tr>
        <td>name</td>
        <td>{name}</td>
        </tr>
        <tr>
        <td>type</td>
        <td>{type}</td>
        </tr>
        <tr>
        <td>filename</td>
        <td>{filename}</td>
        </tr>
        </table>""".format(
            name=infos["name"],
            type=infos["type"],
            filename=infos["filename"].replace("<", "&lt").replace(">", "&gt"),
        )

        # Equations and residuals
        resids = self.residuals()
        equations = self.equations.copy()
        # Create definitions equations and append to equations dictionary
        definitions = self.definitions
        tmp = []
        for deftype in definitions:
            tmp.append(deftype + " = " + definitions[deftype])

        definitions = {"definitions": tmp}
        equations.update(definitions)

        variables = sum([e for k, e in self.symbols.items() if k != "parameters"], [])
        table = '<tr><td><b>Type</b></td><td style="width:80%"><b>Equation</b></td><td><b>Residual</b></td></tr>\n'

        for eq_type in equations:

            eq_lines = []
            for i in range(len(equations[eq_type])):
                eq = equations[eq_type][i]
                # if eq_type in ('expectation','direct_response'):
                #     vals = ''
                if eq_type not in ("arbitrage", "transition", "arbitrage_exp"):
                    vals = ""
                else:
                    val = resids[eq_type][i]
                    if abs(val) > 1e-8:
                        vals = '<span style="color: red;">{:.4f}</span>'.format(val)
                    else:
                        vals = "{:.3f}".format(val)
                if "⟂" in eq:
                    # keep only lhs for now
                    eq, comp = str.split(eq, "⟂")
                if "|" in eq:
                    # keep only lhs for now
                    eq, comp = str.split(eq, "|")
                lat = eq2tex(variables, eq)
                lat = "${}$".format(lat)
                line = [lat, vals]
                h = eq_type if i == 0 else ""
                fmt_line = "<tr><td>{}</td><td>{}</td><td>{}</td></tr>".format(h, *line)
                #         print(fmt_line)
                eq_lines.append(fmt_line)
            table += str.join("\n", eq_lines)
        table = "<table>{}</table>".format(table)

        return table_infos + table

    @property
    def x_bounds(self):

        if "controls_ub" in self.functions:
            fun_lb = self.functions["controls_lb"]
            fun_ub = self.functions["controls_ub"]
            return [fun_lb, fun_ub]
        elif "arbitrage_ub" in self.functions:
            fun_lb = self.functions["arbitrage_lb"]
            fun_ub = self.functions["arbitrage_ub"]
            return [fun_lb, fun_ub]
        else:
            return None

    def residuals(self, calib=None):

        from dolo.algos.steady_state import residuals

        return residuals(self, calib)

    def eval_formula(self, expr, dataframe=None, calib=None):

        from dolo.compiler.eval_formula import eval_formula

        if calib is None:
            calib = self.calibration
        return eval_formula(expr, dataframe=dataframe, context=calib)

`recalibrate(*pargs, **kwargs)` ¶

Return a new, independent model with the given calibration overrides.

Models are immutable (dyno semantics): the calibration of self is left untouched. Assign the returned model instead of mutating in place::

model = model.recalibrate(delta=0.5)

Source code in src/dolo/compiler/model.py

def recalibrate(self, *pargs, **kwargs):
    """Return a new, independent model with the given calibration
    overrides.

    Models are immutable (dyno semantics): the calibration of ``self``
    is left untouched.  Assign the returned model instead of mutating
    in place::

        model = model.recalibrate(delta=0.5)
    """
    if len(pargs) == 1:
        kwargs = dict(pargs[0], **kwargs)
    fresh = copy.copy(self)
    fresh.data = copy.deepcopy(self.data)
    fresh.data["calibration"].update(kwargs)
    fresh.set_changed()
    return fresh

`decode_complementarity(comp, control)` ¶

comp can be either:¶

None
"a<=expr" where a is a controls
"expr<=a" where a is a control
"expr1<=a<=expr2"

Source code in src/dolo/compiler/model.py

def decode_complementarity(comp, control):
    """
    # comp can be either:
    - None
    - "a<=expr" where a is a controls
    - "expr<=a" where a is a control
    - "expr1<=a<=expr2"
    """

    try:
        res = regex.match(comp).groups()
    except:
        raise Exception("Unable to parse complementarity condition '{}'".format(comp))

    res = [r.strip() for r in res]
    if res[1] != control:
        msg = "Complementarity condition '{}' incorrect. Expected {} instead of {}.".format(
            comp, control, res[1]
        )
        raise Exception(msg)

    return [res[0], res[2]]

Importing models¶

`dyno_import(fname, filename=None, check=True)` ¶

Import a .dyno model file using the dynspec-based importer (ModelV2).

Source code in src/dolo/compiler/model_import.py

def dyno_import(fname, filename=None, check=True):
    """Import a ``.dyno`` model file using the dynspec-based importer (ModelV2)."""
    from dolo.compiler.model_v2 import ModelV2

    txt = read_file_or_url(fname)
    if filename is None:
        filename = str(fname)
    return ModelV2.from_dyno(txt, filename=filename, check=check)

`yaml_import_v2(fname, filename=None, check=True)` ¶

Import a dolo YAML model using the dynspec-based importer (ModelV2).

Source code in src/dolo/compiler/model_import.py

def yaml_import_v2(fname, filename=None, check=True):
    """Import a dolo YAML model using the dynspec-based importer (ModelV2)."""
    from dolo.compiler.model_v2 import yaml_compose, ModelV2

    txt = read_file_or_url(fname)
    data = yaml_compose(txt)
    if filename is None:
        filename = str(fname)
    return ModelV2.from_yaml(data, filename=filename, check=check)

Steady state and solution algorithms¶

Time Iteration Algorithm

`euler_residuals_tau(dmodel, controls, dr)` ¶

Universal Euler residuals written entirely in terms of GArray and Tau.

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

def euler_residuals_tau(
    dmodel: DiscretizedModel, controls: GArray, dr: DecisionRule
) -> GArray:
    """Universal Euler residuals written entirely in terms of GArray and Tau."""
    res = GArray.zeros_like(controls)
    for i_m, m_val, s_endo in dmodel.regimes():
        xm = controls[i_m]
        for prob, S in dmodel.tau.slice_transitions(i_m, m_val, s_endo, xm):
            XM = dr(S)
            rr = dmodel.arbitrage(m_val, s_endo, xm, S.exo, S.endo, XM)
            res[i_m] += prob * rr
    return res

`residuals_simple(f, g, s, x, dr, dprocess, parms)` ¶

Legacy residuals computation retained for backward compatibility.

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

def residuals_simple(f, g, s, x, dr, dprocess, parms):
    """Legacy residuals computation retained for backward compatibility."""
    N = s.shape[0]
    res = numpy.zeros_like(x)

    for i_ms in range(dprocess.n_nodes):
        m = numpy.tile(dprocess.node(i_ms), (N, 1))
        xm = x[i_ms, :, :]

        for I_ms in range(dprocess.n_inodes(i_ms)):
            M = numpy.tile(dprocess.inode(i_ms, I_ms), (N, 1))
            prob = dprocess.iweight(i_ms, I_ms)
            S = g(m, s, xm, M, parms)
            XM = dr.eval_ijs(i_ms, I_ms, S)
            rr = f(m, s, xm, M, S, XM, parms)
            res[i_ms, :, :] += prob * rr

    return res

`time_iteration(model, *, engine='auto', backend=None, dr0=None, tol_eps=1e-08, tol_eta=1e-06, maxit=500, inner=None, accelerate=False, interpolation='cubic', verbose=True, details=True, trace=False, ignore_constraints=False, hook=None, dprocess=None, **kwargs)` ¶

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,
    )

Newton-Krylov Policy Solver

`Operator` ¶

Bases: LinearOperator

Special Linear Operator representing the policy Jacobian L = -F_1^{-1} F_2

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

class Operator(LinearOperator):
    """Special Linear Operator representing the policy Jacobian L = -F_1^{-1} F_2"""

    def __init__(self, M_ij, S_ij, dumdr):
        self.M_ij = M_ij
        self.S_ij = S_ij
        self.n_m = M_ij.shape[0]
        self.N = M_ij.shape[2]
        self.n_x = M_ij.shape[3]
        self.dumdr = dumdr
        self.dtype = numpy.dtype("float64")
        self.counter = 0
        self.addid = False
        nn = self.n_m * self.N * self.n_x
        super().__init__(dtype=self.dtype, shape=(nn, nn))
        self.ndim = 2

    def _matvec(self, x):
        self.counter += 1
        xx = x.reshape((self.n_m, self.N, self.n_x))
        yy = self.apply(xx)
        if self.addid:
            yy = xx - yy  # (I - L) operator
        return yy.ravel()

    def apply(self, pi, inplace=False):
        M_ij = self.M_ij
        S_ij = self.S_ij
        n_m = self.n_m
        N = self.N
        n_x = self.n_x
        dumdr = self.dumdr
        if not inplace:
            pi = pi.copy()
        return d_filt_dx(pi, M_ij, S_ij, n_m, N, n_x, dumdr)

    def as_matrix(self):
        arg = numpy.zeros((self.n_m, self.N, self.n_x))
        larg = arg.ravel()
        N = len(larg)
        J = numpy.zeros((N, N))
        for i in range(N):
            if i > 0:
                larg[i - 1] = 0.0
            larg[i] = 1.0
            J[:, i] = self.apply(arg).ravel()
        return J

`SparseTensor` ¶

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

class SparseTensor:
    def __init__(self, indices, values, shape):
        """
        indices: N x d integers
        values:  N     floats
        """
        indices = numpy.array(indices, dtype=int)
        values = numpy.array(values, dtype=float)
        self.indices = indices
        self.values = values
        self.shape = shape

    @property
    def lin_indices(self):
        minds = tuple(self.indices.T)
        return numpy.ravel_multi_index(minds, self.shape)

    def reshape(self, shape):
        assert prod(shape) == prod(self.shape)
        new_indices = numpy.unravel_index(self.lin_indices, shape)
        new_indices = numpy.concatenate([a[:, None] for a in new_indices], axis=1)
        return SparseTensor(new_indices, self.values, shape)

    def as_spmatrix(self, dims=None):
        import scipy.sparse

        if dims is None:
            pp = prod(self.shape)
            P = int(sqrt(pp))
            Q = P
            assert pp == P**2
            dims = (P, Q)
        else:
            assert len(dims) == 2
        spmat = self.reshape(dims)
        data = spmat.values
        inds = spmat.indices
        return scipy.sparse.coo_matrix((data, (inds[:, 0], inds[:, 1])))

    def as_array(self, dims=None):
        return self.as_spmatrix(dims).toarray()

`init(indices, values, shape)` ¶

indices: N x d integers values: N floats

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

def __init__(self, indices, values, shape):
    """
    indices: N x d integers
    values:  N     floats
    """
    indices = numpy.array(indices, dtype=int)
    values = numpy.array(values, dtype=float)
    self.indices = indices
    self.values = values
    self.shape = shape

`golden_section_search(f_eval, alpha_min=0.0, alpha_max=1.0, maxiter=8, tol=0.0001)` ¶

Golden-section line search for step size alpha in [alpha_min, alpha_max].

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

def golden_section_search(f_eval, alpha_min=0.0, alpha_max=1.0, maxiter=8, tol=1e-4):
    """Golden-section line search for step size alpha in [alpha_min, alpha_max]."""
    gr = (sqrt(5.0) - 1.0) / 2.0  # ~0.6180339887
    a = alpha_min
    b = alpha_max
    c = b - gr * (b - a)
    d = a + gr * (b - a)

    fc = f_eval(c)
    fd = f_eval(d)

    best_alpha = c if fc < fd else d
    best_f = min(fc, fd)

    for _ in range(maxiter):
        if fc < fd:
            b = d
            d = c
            fd = fc
            c = b - gr * (b - a)
            fc = f_eval(c)
        else:
            a = c
            c = d
            fc = fd
            d = a + gr * (b - a)
            fd = f_eval(d)

        if fc < best_f:
            best_f = fc
            best_alpha = c
        if fd < best_f:
            best_f = fd
            best_alpha = d
        if abs(b - a) < tol:
            break

    return best_alpha, best_f

`invert_jac(res, dres, jres, fut_S, dumdr, tol=1e-10, maxit=1000, verbose=False)` ¶

Solve (F_1 + F_2) tot = res via preconditioned Neumann series (I - L) tot = F_1^{-1} res.

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

def invert_jac(res, dres, jres, fut_S, dumdr, tol=1e-10, maxit=1000, verbose=False):
    """Solve (F_1 + F_2) tot = res via preconditioned Neumann series (I - L) tot = F_1^{-1} res."""
    ddx = solve_gu(dres.copy(), res.copy())
    return neumann_solve(ddx, jres, fut_S, dumdr, tol=tol, maxit=maxit, verbose=verbose)

`neumann_solve(ddx, jres, fut_S, dumdr, tol=1e-08, maxit=50, verbose=False)` ¶

Solve (I - L) tot = ddx using Neumann series expansion: sum_{k=0}^maxit L^k ddx.

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

def neumann_solve(ddx, jres, fut_S, dumdr, tol=1e-8, maxit=50, verbose=False):
    """Solve (I - L) tot = ddx using Neumann series expansion: sum_{k=0}^maxit L^k ddx."""
    n_m = ddx.shape[0]
    N = ddx.shape[1]
    n_x = ddx.shape[2]

    tot = ddx.copy()
    du = ddx.copy()
    err_0 = abs(du).max()
    lam = -1.0
    lam_max = -1.0
    nn = 0

    if verbose:
        print("Starting Neumann series")
    for nn in range(maxit):
        du = d_filt_dx(du, jres, fut_S, n_m, N, n_x, dumdr)
        err = abs(du).max()
        lam = err / err_0 if err_0 > 0 else 0.0
        lam_max = max(lam_max, lam)
        if verbose:
            print(f"- {err} | {lam} | {lam_max}")
        tot += du
        err_0 = err
        if err < tol:
            break

    return tot, nn, lam

`newton(model, *, engine='auto', backend=None, device=None, dr0=None, tol_eps=1e-08, tol_eta=1e-06, maxit=30, linsolve='gmres', linesearch='backtracking', interpolation='cubic', verbose=True, details=True, trace=False, dprocess=None, ignore_constraints=False, compute_radius=False, **kwargs)` ¶

Finds a global solution for model using full-system Newton-Krylov solver.

Parameters¶

model : Model model to be solved 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 (30) maximum number of outer Newton iterations linsolve : str or dict ("gmres") linear solver for the preconditioned system (I - L) dx = r. Options: "gmres", "neumann", or a dict(method=..., maxit=..., tol=...) linesearch : str or dict ("backtracking") step size stabilization. Options: "backtracking", "golden_section", False, None, or a dict(method=..., max_backsteps=...) 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/newton.py

def newton(
    model: Model,
    *,
    engine: str = "auto",
    backend: Optional[str] = None,
    device: Optional[Any] = None,
    dr0: DecisionRule = None,
    tol_eps: float = 1e-8,
    tol_eta: float = 1e-6,
    maxit: int = 30,
    linsolve="gmres",
    linesearch="backtracking",
    interpolation: str = "cubic",
    verbose: bool = True,
    details: bool = True,
    trace: bool = False,
    dprocess=None,
    ignore_constraints: bool = False,
    compute_radius: bool = False,
    **kwargs,
) -> NewtonResult:
    """Finds a global solution for ``model`` using full-system Newton-Krylov solver.

    Parameters
    ----------
    model : Model
        model to be solved
    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 (30)
        maximum number of outer Newton iterations
    linsolve : str or dict ("gmres")
        linear solver for the preconditioned system (I - L) dx = r.
        Options: "gmres", "neumann", or a dict(method=..., maxit=..., tol=...)
    linesearch : str or dict ("backtracking")
        step size stabilization. Options: "backtracking", "golden_section", False, None,
        or a dict(method=..., max_backsteps=...)
    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_eps = kwargs.pop("tol")
    if "interp_method" in kwargs:
        interpolation = kwargs.pop("interp_method")

    complementarities = kwargs.pop("complementarities", None)
    invmethod = kwargs.pop("invmethod", None)
    smaxit = kwargs.pop("smaxit", None)
    linear_maxit = kwargs.pop("linear_maxit", None)
    max_backsteps = kwargs.pop("max_backsteps", None)
    maxbsteps = kwargs.pop("maxbsteps", None)
    mu = kwargs.pop("mu", 2)
    method = kwargs.pop("method", "jac")

    if complementarities is not None:
        pass
    else:
        complementarities = not ignore_constraints

    # Linear solver config
    if isinstance(linsolve, str):
        linsolve_cfg = {
            "method": "neumann" if linsolve == "iti" else linsolve,
            "maxit": 50,
            "tol": 1e-8,
            "precondition": True,
        }
    elif isinstance(linsolve, dict):
        linsolve_cfg = {
            "method": "gmres",
            "maxit": 50,
            "tol": 1e-8,
            "precondition": True,
            **linsolve,
        }
        if linsolve_cfg["method"] == "iti":
            linsolve_cfg["method"] = "neumann"
    else:
        linsolve_cfg = {
            "method": "gmres",
            "maxit": 50,
            "tol": 1e-8,
            "precondition": True,
        }

    if invmethod is not None:
        linsolve_cfg["method"] = "neumann" if invmethod == "iti" else invmethod
    if smaxit is not None:
        linsolve_cfg["maxit"] = smaxit
    if linear_maxit is not None:
        linsolve_cfg["maxit"] = linear_maxit

    # Line search config
    if linesearch is False or linesearch is None or linesearch == "none":
        linesearch_cfg = {"method": None, "max_backsteps": 0}
    elif isinstance(linesearch, str):
        linesearch_cfg = {"method": linesearch, "max_backsteps": 10}
    elif isinstance(linesearch, dict):
        linesearch_cfg = {"method": "backtracking", "max_backsteps": 10, **linesearch}
    else:
        linesearch_cfg = {"method": "backtracking", "max_backsteps": 10}

    if max_backsteps is not None:
        linesearch_cfg["max_backsteps"] = max_backsteps
    elif maxbsteps is not None:
        linesearch_cfg["max_backsteps"] = maxbsteps

    itprint = IterationsPrinter(
        ("N", int),
        ("f_x", float),
        ("d_x", float),
        ("Time_residuals", float),
        ("Time_inversion", float),
        ("Time_search", float),
        ("Lambda_0", float),
        ("N_invert", int),
        ("N_search", int),
        verbose=verbose,
    )
    itprint.print_header("Start Newton Iterations.")

    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

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

    f = dmodel.f
    g = dmodel.g
    x_lb = dmodel.model.functions.get("arbitrage_lb", None)
    x_ub = dmodel.model.functions.get("arbitrage_ub", None)

    parms = dmodel.parms

    n_m = max(dprocess.n_nodes, 1)
    n_s = len(dmodel.model.symbols["states"])

    eng = _resolve_engine(dmodel, engine, backend)
    use_jax = eng in ("jax", "gpu", "cpu") and n_m == 1
    if use_jax:
        target_dev = device
        if target_dev is None:
            if eng == "gpu":
                import jax

                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":
                import jax

                target_dev = jax.devices("cpu")[0]
    else:
        target_dev = None

    if interpolation in ("cubic", "linear"):
        ddr = DecisionRule(
            exo_grid, endo_grid, dprocess=dprocess, interp_method=interpolation
        )
        ddr_filt = DecisionRule(
            exo_grid, endo_grid, dprocess=dprocess, interp_method=interpolation
        )
    else:
        raise Exception("Unsupported interpolation method.")

    s = endo_grid.nodes
    N = s.shape[0]
    n_x = len(dmodel.model.symbols["controls"])
    x0 = (
        dmodel.model.calibration["controls"][
            None,
            None,
        ]
        .repeat(n_m, axis=0)
        .repeat(N, axis=1)
    )

    if dr0 is not None:
        for i_m in range(n_m):
            x0[i_m, :, :] = dr0(i_m, s)
    ddr.set_values(x0)

    if x_lb is not None and complementarities:
        lb = x0.copy()
        ub = x0.copy()
        for i_m in range(n_m):
            m = dprocess.node(i_m)
            lb[i_m, :] = x_lb(m, s, parms)
            ub[i_m, :] = x_ub(m, s, parms)
    else:
        complementarities = False
        lb = None
        ub = None

    x = x0.copy()
    ddr.set_values(x)

    n_im = dprocess.n_inodes(0)
    jres = numpy.zeros((n_m, n_im, N, n_x, n_x))
    S_ij = numpy.zeros((n_m, n_im, N, n_s))

    trace_details = [] if trace else None

    err_0 = 10.0
    err_2 = 10.0

    for it in range(maxit):
        if use_jax and not complementarities:
            import jax
            import jax.numpy as jnp

            t1 = time.time()
            ddr.set_values(x)

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

            # 1. Residuals and Jacobian
            res, dres = F_and_dF_1(
                dmodel,
                x,
                dr=ddr,
                engine=eng,
                interp_method=interpolation,
                device=target_dev,
            )
            res_arr = numpy.asarray(
                res.data[0]
                if isinstance(res, GArray)
                else (res[0] if res.ndim == 3 else res)
            )
            dres_arr = numpy.asarray(
                dres.data[0]
                if isinstance(dres, GArray)
                else (dres[0] if dres.ndim == 4 else dres)
            )

            err_0 = abs(res_arr).max()
            if err_0 < tol_eps:
                t2 = time.time()
                t3 = t2
                t4 = t2
                itprint.print_iteration(
                    N=it,
                    f_x=err_0,
                    d_x=0.0,
                    Time_residuals=t2 - t1,
                    Time_inversion=0.0,
                    Time_search=0.0,
                    Lambda_0=0.0,
                    N_invert=0,
                    N_search=0,
                )
                break

            t2 = time.time()

            # 2. Block-diagonal solve: ddx = dres^{-1} * res
            J = jnp.asarray(dres_arr)
            R = jnp.asarray(res_arr)
            ddx = numpy.asarray(jnp.linalg.solve(J, R[..., None])[..., 0])

            # 3. Policy linear operator dF_2 and preconditioned inversion
            solver_method = linsolve_cfg.get("method", "gmres")
            solver_maxit = linsolve_cfg.get("maxit", 50)
            solver_tol = linsolve_cfg.get("tol", 1e-8)
            ttol = max(err_0 / 100, solver_tol)

            L_op = dF_2(
                dmodel,
                x,
                dr=ddr,
                engine=eng,
                interp_method=interpolation,
                device=target_dev,
            )
            L_prec = (-1.0) * L_op.preconditioned(J)
            tot_2d = L_prec.solve(
                ddx, linsolve=solver_method, tol=ttol, maxit=solver_maxit
            )
            tot = numpy.asarray(tot_2d).reshape(x.shape)
            L_last = L_prec
            nn = solver_maxit
            lam0 = 0.01

            t3 = time.time()

            # 4. Line search
            ls_method = linesearch_cfg.get("method", "backtracking")
            ls_maxsteps = linesearch_cfg.get("max_backsteps", 10)

            def eval_res_norm_jax(cand_x):
                c_res = F(
                    dmodel,
                    cand_x,
                    dr=ddr,
                    engine=eng,
                    interp_method=interpolation,
                    device=target_dev,
                )
                c_arr = numpy.asarray(
                    c_res.data[0]
                    if isinstance(c_res, GArray)
                    else (c_res[0] if c_res.ndim == 3 else c_res)
                )
                return abs(c_arr).max()

            if ls_method == "golden_section":
                cand_1 = eval_res_norm_jax(x - tot)
                if cand_1 < err_0:
                    best_alpha = 1.0
                    i_bckstps = 0
                else:
                    best_alpha, _ = golden_section_search(
                        lambda a: eval_res_norm_jax(x - tot * a),
                        alpha_min=0.0,
                        alpha_max=1.0,
                        maxiter=ls_maxsteps,
                    )
                    i_bckstps = ls_maxsteps
                new_x = x - tot * best_alpha

            elif ls_method == "backtracking":
                steps = 0.5 ** numpy.arange(max(ls_maxsteps, 1))
                new_x = x - tot
                i_bckstps = 0
                for i_bckstps, lam in enumerate(steps):
                    new_x = x - tot * lam
                    new_err = eval_res_norm_jax(new_x)
                    if new_err < err_0:
                        break
            else:
                new_x = x - tot
                i_bckstps = 0

            err_2 = abs(tot).max()
            t4 = time.time()

            itprint.print_iteration(
                N=it,
                f_x=err_0,
                d_x=err_2,
                Time_residuals=t2 - t1,
                Time_inversion=t3 - t2,
                Time_search=t4 - t3,
                Lambda_0=lam0,
                N_invert=nn,
                N_search=i_bckstps,
            )

            x = new_x
            if err_0 < tol_eps or err_2 < tol_eta:
                break

        else:
            jres[...] = 0.0
            S_ij[...] = 0.0

            t1 = time.time()
            ddr.set_values(x)

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

            sh_x = x.shape
            ff = SerialDifferentiableFunction(
                lambda u: euler_residuals(
                    f,
                    g,
                    s,
                    u.reshape(sh_x),
                    ddr,
                    dprocess,
                    parms,
                    diff=False,
                    with_jres=False,
                    set_dr=False,
                ).reshape((-1, sh_x[2]))
            )
            res, dres = ff(x.reshape((-1, sh_x[2])))
            res = res.reshape(sh_x)
            dres = dres.reshape((sh_x[0], sh_x[1], sh_x[2], sh_x[2]))
            junk, jres, fut_S = euler_residuals(
                f,
                g,
                s,
                x,
                ddr,
                dprocess,
                parms,
                diff=False,
                with_jres=True,
                set_dr=False,
                jres=jres,
                S_ij=S_ij,
            )

            if complementarities:
                res, dres, jres = smooth(res, dres, jres, x - lb)
                res[...] *= -1
                dres[...] *= -1
                jres[...] *= -1
                res, dres, jres = smooth(res, dres, jres, ub - x, pos=-1.0)
                res[...] *= -1
                dres[...] *= -1
                jres[...] *= -1

            err_0 = abs(res).max()

            jres[...] *= -1.0

            for i_m in range(n_m):
                for j_m in range(n_im):
                    M = jres[i_m, j_m, :, :, :]
                    X = dres[i_m, :, :, :].copy()
                    sol = solve_tensor(X, M)

            t2 = time.time()

            # Inversion of (I - L) dx = F_1^{-1} res
            solver_method = linsolve_cfg.get("method", "gmres")
            solver_maxit = linsolve_cfg.get("maxit", 50)
            solver_tol = linsolve_cfg.get("tol", 1e-8)

            if solver_method == "gmres":
                ddx = solve_gu(dres.copy(), res.copy())
                L = Operator(jres, fut_S, ddr_filt)
                n0 = L.counter
                L.addid = True
                ttol = max(err_0 / 100, solver_tol)
                try:
                    sol, info = scipy.sparse.linalg.gmres(
                        L, ddx.ravel(), rtol=ttol, atol=0.0, maxiter=solver_maxit
                    )
                except TypeError:
                    sol, info = scipy.sparse.linalg.gmres(
                        L, ddx.ravel(), tol=ttol, maxiter=solver_maxit
                    )
                lam0 = 0.01
                nn = L.counter - n0
                tot = sol.reshape(ddx.shape)
            else:
                # Neumann series
                tot, nn, lam0 = invert_jac(
                    res,
                    dres,
                    jres,
                    fut_S,
                    ddr_filt,
                    tol=solver_tol,
                    maxit=solver_maxit,
                    verbose=(verbose == "full"),
                )

            t3 = time.time()

            # Step search / Line search
            ls_method = linesearch_cfg.get("method", "backtracking")
            ls_maxsteps = linesearch_cfg.get("max_backsteps", 10)

            def eval_residual_norm(cand_x):
                cand_res = euler_residuals(
                    f, g, s, cand_x, ddr, dprocess, parms, diff=False, set_dr=True
                )
                if complementarities:
                    cand_res = smooth_nodiff(cand_res, cand_x - lb)
                    cand_res = smooth_nodiff(-cand_res, ub - cand_x)
                return abs(cand_res).max()

            if ls_method == "golden_section":
                # Check full step first
                cand_1 = eval_residual_norm(x - tot)
                if cand_1 < err_0:
                    best_alpha = 1.0
                    i_bckstps = 0
                else:
                    best_alpha, _ = golden_section_search(
                        lambda a: eval_residual_norm(x - tot * a),
                        alpha_min=0.0,
                        alpha_max=1.0,
                        maxiter=ls_maxsteps,
                    )
                    i_bckstps = ls_maxsteps
                new_x = x - tot * best_alpha

            elif ls_method == "backtracking":
                steps = 0.5 ** numpy.arange(max(ls_maxsteps, 1))
                new_x = x - tot
                for i_bckstps, lam in enumerate(steps):
                    new_x = x - tot * lam
                    new_err = eval_residual_norm(new_x)
                    if new_err < err_0:
                        break
            else:
                # No line search (full Newton step)
                new_x = x - tot
                i_bckstps = 0

            err_2 = abs(tot).max()
            t4 = time.time()

            itprint.print_iteration(
                N=it,
                f_x=err_0,
                d_x=err_2,
                Time_residuals=t2 - t1,
                Time_inversion=t3 - t2,
                Time_search=t4 - t3,
                Lambda_0=lam0,
                N_invert=nn,
                N_search=i_bckstps,
            )

            if err_0 < tol_eps:
                break

            x = new_x

    ddr.set_values(x)
    if trace:
        trace_details.append({"dr": copy.deepcopy(ddr)})

    itprint.print_finished()

    if not details:
        return ddr
    else:
        if use_jax and not complementarities:
            L = L_last if "L_last" in locals() else None
            if compute_radius and L is not None:
                try:
                    lam_eig = scipy.sparse.linalg.eigs(L, k=1, return_eigenvectors=False)
                    lam = abs(lam_eig[0])
                except Exception:
                    lam = numpy.nan
            else:
                lam = numpy.nan
        else:
            L = Operator(jres, fut_S, ddr_filt)
            if compute_radius:
                try:
                    lam_eig = scipy.sparse.linalg.eigs(L, k=1, return_eigenvectors=False)
                    lam = abs(lam_eig[0])
                except Exception:
                    lam = numpy.nan
            else:
                lam = numpy.nan

        return NewtonResult(
            ddr,
            it + 1,
            err_0,
            err_2,
            err_0 < tol_eps,
            complementarities,
            lam,
            trace_details,
            L,
        )

`qzordered(A, B, crit=1.0)` ¶

Eigenvalues bigger than crit are sorted in the top-left.

Source code in src/dolo/numeric/extern/qz.py

def qzordered(A, B, crit=1.0):
    "Eigenvalues bigger than crit are sorted in the top-left."

    TOL = 1e-10

    def select(alpha, beta):
        return alpha**2 > crit * beta**2

    [S, T, alpha, beta, U, V] = ordqz(A, B, output="real", sort=select)

    eigval = abs(numpy.diag(S) / numpy.diag(T))

    return [S, T, U, V, eigval]

Numeric Core Abstractions¶

The numeric core provides the grid, transition kernel, and discrete field abstractions:

DiscretizedModel: The primary representation of a discretized economic model.

Following Dolo.jl: A DiscretizedModel bundles: - The continuous/symbolic Model - The spatial Grid - The transition operator Tau - Methods for initial guess, bounds, arbitrage, and transition matrix generation.

`DiscretizedModel` ¶

Discretized economic model.

Encapsulates: - model: Symbolic Model / ModelV2 - grid: Spatial Grid (ProductGrid or CartesianGrid) - tau: Tau transition operator - dprocess: Underlying discretized process

Source code in src/dolo/numeric/discretized_model.py

class DiscretizedModel:
    """Discretized economic model.

    Encapsulates:
    - `model`: Symbolic Model / ModelV2
    - `grid`: Spatial Grid (ProductGrid or CartesianGrid)
    - `tau`: Tau transition operator
    - `dprocess`: Underlying discretized process
    """

    def __init__(
        self,
        model: Any,
        grid_options: Optional[Dict[str, Any]] = None,
        dprocess_options: Optional[Dict[str, Any]] = None,
    ):
        self.model = model
        dprocess_options = dprocess_options or {}

        # Discretize model domain and process
        self.grid, self.dprocess = model.discretize(grid_options, dprocess_options)
        self.parms = model.calibration["parameters"]
        self.f = model.functions["arbitrage"]
        self.g = model.functions["transition"]

        # Tau transition operator
        self.tau = Tau(self.g, self.dprocess, self.parms)

        # Grid characteristics
        self.is_mc = not isinstance(self.grid.grids[0], EmptyGrid)
        self.n_regimes = self.dprocess.n_nodes
        self.endo_grid = self.grid.grids[1]
        self.endo_nodes = self.endo_grid.nodes  # shape (N, n_s)
        self.N_endo = self.endo_nodes.shape[0]
        self.N_total = self.n_regimes * self.N_endo

        # Mathematical Space representation
        from dolo.numeric.spaces import (
            DiscreteSpace,
            EmptySpace,
            ProductSpace,
            CartesianSpace,
        )

        exo_space = (
            DiscreteSpace(self.dprocess.grid.nodes) if self.is_mc else EmptySpace()
        )
        if hasattr(model, "domain") and model.domain is not None:
            endo_space = model.domain
        else:
            endo_space = CartesianSpace(min=self.endo_grid.min, max=self.endo_grid.max)
        self.space = ProductSpace([exo_space, endo_space], names=["exo", "endo"])

        # Approximator
        from dolo.numeric.approximators import get_approximator

        interp_method = (grid_options or {}).get("interp_method", "cubic")
        try:
            self.approximator = get_approximator(
                self.grid, interp_method=interp_method, space=self.space
            )
        except Exception:
            self.approximator = None

    def decision_rule(
        self,
        values: Optional[Union[GArray, np.ndarray]] = None,
        interp_method: str = "cubic",
    ) -> Any:
        """Create a DecisionRule bound to this discretized model."""
        from dolo.numeric.decision_rule import DecisionRule

        dr = DecisionRule(
            self.grid.grids[0],
            self.grid.grids[1],
            interp_method=interp_method,
            dprocess=self.dprocess,
        )
        if values is not None:
            dr.fit(values)
        return dr

    def regimes(self) -> Iterator[Tuple[int, np.ndarray, np.ndarray]]:
        """Yields (regime_index, m_value, endo_nodes) for each discrete regime."""
        for i in range(self.n_regimes):
            m_val = self.dprocess.node(i)
            yield i, m_val, self.endo_nodes

    def arbitrage(
        self,
        m: np.ndarray,
        s: np.ndarray,
        x: np.ndarray,
        M: np.ndarray,
        S: np.ndarray,
        X: np.ndarray,
    ) -> np.ndarray:
        """Arbitrage equation evaluation wrapper."""
        N = s.shape[0] if s.ndim == 2 else 1
        m_tiled = np.tile(m, (N, 1)) if m.ndim == 1 and N > 1 else m
        M_tiled = np.tile(M, (N, 1)) if M.ndim == 1 and N > 1 else M
        return self.f(m_tiled, s, x, M_tiled, S, X, self.parms)

    def initial_guess(self) -> GArray:
        """Initial GArray of controls from model calibration."""
        x0 = self.model.calibration["controls"]
        n_x = len(x0)
        controls = GArray.zeros(self.grid, n_x, self.n_regimes)
        controls.data[:, :, :] = x0[None, None, :]
        return controls

    def lower_bounds(self) -> Optional[GArray]:
        """Lower bounds on controls as GArray, if specified."""
        if "arbitrage_lb" not in self.model.functions:
            return None
        lb_fun = self.model.functions["arbitrage_lb"]
        controls_lb = GArray.zeros(
            self.grid, len(self.model.symbols["controls"]), self.n_regimes
        )
        N = self.N_endo
        p = np.repeat(self.parms[None, :], N, axis=0)

        for i, m_val, s_endo in self.regimes():
            m = np.repeat(m_val[None, :], N, axis=0)
            controls_lb[i] = lb_fun(m, s_endo, p)
        return controls_lb

    def upper_bounds(self) -> Optional[GArray]:
        """Upper bounds on controls as GArray, if specified."""
        if "arbitrage_ub" not in self.model.functions:
            return None
        ub_fun = self.model.functions["arbitrage_ub"]
        controls_ub = GArray.zeros(
            self.grid, len(self.model.symbols["controls"]), self.n_regimes
        )
        N = self.N_endo
        p = np.repeat(self.parms[None, :], N, axis=0)

        for i, m_val, s_endo in self.regimes():
            m = np.repeat(m_val[None, :], N, axis=0)
            controls_ub[i] = ub_fun(m, s_endo, p)
        return controls_ub

    def transition_matrix(self, controls: GArray) -> np.ndarray:
        """Compute the full Markov transition matrix P on the grid nodes using tau_fit.

        Parameters
        ----------
        controls : GArray
            Policy controls defined on the grid.

        Returns
        -------
        P : np.ndarray, shape (N_total, N_total)
            Stochastic matrix where P[i, j] is the transition probability from node i to node j.
        """
        P = np.zeros((self.N_total, self.N_total), dtype=float)
        N_e = self.N_endo

        for i_m, m_val, s_endo in self.regimes():
            xm = controls[i_m]
            row_offset = i_m * N_e

            for j_target, prob, w_corner, idx_corner in self.tau.tau_fit(
                i_m, m_val, s_endo, xm, self.endo_grid
            ):
                col_offset = j_target * N_e
                for n in range(N_e):
                    row = row_offset + n
                    col = col_offset + idx_corner[n]
                    P[row, col] += prob * w_corner[n]

        return P

    def ergodic_distribution(self, controls: GArray) -> GArray:
        """Compute the ergodic stationary distribution on the grid nodes.

        Parameters
        ----------
        controls : GArray
            Policy controls defined on the grid.

        Returns
        -------
        mu : GArray
            Stationary probability distribution across all grid nodes.
        """
        P = self.transition_matrix(controls)
        # Solve mu P = mu  <=>  (P^T - I) mu = 0 with sum(mu) = 1
        A = P.T - np.eye(self.N_total)
        A[-1, :] = 1.0
        b = np.zeros(self.N_total)
        b[-1] = 1.0
        mu_flat = np.linalg.solve(A, b)
        mu_flat = np.maximum(mu_flat, 0.0)
        mu_flat /= mu_flat.sum()

        mu = GArray.zeros(self.grid, 1, self.n_regimes)
        return mu.unravel(mu_flat)

`arbitrage(m, s, x, M, S, X)` ¶

Arbitrage equation evaluation wrapper.

Source code in src/dolo/numeric/discretized_model.py

def arbitrage(
    self,
    m: np.ndarray,
    s: np.ndarray,
    x: np.ndarray,
    M: np.ndarray,
    S: np.ndarray,
    X: np.ndarray,
) -> np.ndarray:
    """Arbitrage equation evaluation wrapper."""
    N = s.shape[0] if s.ndim == 2 else 1
    m_tiled = np.tile(m, (N, 1)) if m.ndim == 1 and N > 1 else m
    M_tiled = np.tile(M, (N, 1)) if M.ndim == 1 and N > 1 else M
    return self.f(m_tiled, s, x, M_tiled, S, X, self.parms)

`decision_rule(values=None, interp_method='cubic')` ¶

Create a DecisionRule bound to this discretized model.

Source code in src/dolo/numeric/discretized_model.py

def decision_rule(
    self,
    values: Optional[Union[GArray, np.ndarray]] = None,
    interp_method: str = "cubic",
) -> Any:
    """Create a DecisionRule bound to this discretized model."""
    from dolo.numeric.decision_rule import DecisionRule

    dr = DecisionRule(
        self.grid.grids[0],
        self.grid.grids[1],
        interp_method=interp_method,
        dprocess=self.dprocess,
    )
    if values is not None:
        dr.fit(values)
    return dr

`ergodic_distribution(controls)` ¶

Compute the ergodic stationary distribution on the grid nodes.

Parameters¶

controls : GArray Policy controls defined on the grid.

Returns¶

mu : GArray Stationary probability distribution across all grid nodes.

Source code in src/dolo/numeric/discretized_model.py

def ergodic_distribution(self, controls: GArray) -> GArray:
    """Compute the ergodic stationary distribution on the grid nodes.

    Parameters
    ----------
    controls : GArray
        Policy controls defined on the grid.

    Returns
    -------
    mu : GArray
        Stationary probability distribution across all grid nodes.
    """
    P = self.transition_matrix(controls)
    # Solve mu P = mu  <=>  (P^T - I) mu = 0 with sum(mu) = 1
    A = P.T - np.eye(self.N_total)
    A[-1, :] = 1.0
    b = np.zeros(self.N_total)
    b[-1] = 1.0
    mu_flat = np.linalg.solve(A, b)
    mu_flat = np.maximum(mu_flat, 0.0)
    mu_flat /= mu_flat.sum()

    mu = GArray.zeros(self.grid, 1, self.n_regimes)
    return mu.unravel(mu_flat)

`initial_guess()` ¶

Initial GArray of controls from model calibration.

Source code in src/dolo/numeric/discretized_model.py

def initial_guess(self) -> GArray:
    """Initial GArray of controls from model calibration."""
    x0 = self.model.calibration["controls"]
    n_x = len(x0)
    controls = GArray.zeros(self.grid, n_x, self.n_regimes)
    controls.data[:, :, :] = x0[None, None, :]
    return controls

`lower_bounds()` ¶

Lower bounds on controls as GArray, if specified.

Source code in src/dolo/numeric/discretized_model.py

def lower_bounds(self) -> Optional[GArray]:
    """Lower bounds on controls as GArray, if specified."""
    if "arbitrage_lb" not in self.model.functions:
        return None
    lb_fun = self.model.functions["arbitrage_lb"]
    controls_lb = GArray.zeros(
        self.grid, len(self.model.symbols["controls"]), self.n_regimes
    )
    N = self.N_endo
    p = np.repeat(self.parms[None, :], N, axis=0)

    for i, m_val, s_endo in self.regimes():
        m = np.repeat(m_val[None, :], N, axis=0)
        controls_lb[i] = lb_fun(m, s_endo, p)
    return controls_lb

`regimes()` ¶

Yields (regime_index, m_value, endo_nodes) for each discrete regime.

Source code in src/dolo/numeric/discretized_model.py

def regimes(self) -> Iterator[Tuple[int, np.ndarray, np.ndarray]]:
    """Yields (regime_index, m_value, endo_nodes) for each discrete regime."""
    for i in range(self.n_regimes):
        m_val = self.dprocess.node(i)
        yield i, m_val, self.endo_nodes

`transition_matrix(controls)` ¶

Compute the full Markov transition matrix P on the grid nodes using tau_fit.

Parameters¶

controls : GArray Policy controls defined on the grid.

Returns¶

P : np.ndarray, shape (N_total, N_total) Stochastic matrix where P[i, j] is the transition probability from node i to node j.

Source code in src/dolo/numeric/discretized_model.py

def transition_matrix(self, controls: GArray) -> np.ndarray:
    """Compute the full Markov transition matrix P on the grid nodes using tau_fit.

    Parameters
    ----------
    controls : GArray
        Policy controls defined on the grid.

    Returns
    -------
    P : np.ndarray, shape (N_total, N_total)
        Stochastic matrix where P[i, j] is the transition probability from node i to node j.
    """
    P = np.zeros((self.N_total, self.N_total), dtype=float)
    N_e = self.N_endo

    for i_m, m_val, s_endo in self.regimes():
        xm = controls[i_m]
        row_offset = i_m * N_e

        for j_target, prob, w_corner, idx_corner in self.tau.tau_fit(
            i_m, m_val, s_endo, xm, self.endo_grid
        ):
            col_offset = j_target * N_e
            for n in range(N_e):
                row = row_offset + n
                col = col_offset + idx_corner[n]
                P[row, col] += prob * w_corner[n]

    return P

`upper_bounds()` ¶

Upper bounds on controls as GArray, if specified.

Source code in src/dolo/numeric/discretized_model.py

def upper_bounds(self) -> Optional[GArray]:
    """Upper bounds on controls as GArray, if specified."""
    if "arbitrage_ub" not in self.model.functions:
        return None
    ub_fun = self.model.functions["arbitrage_ub"]
    controls_ub = GArray.zeros(
        self.grid, len(self.model.symbols["controls"]), self.n_regimes
    )
    N = self.N_endo
    p = np.repeat(self.parms[None, :], N, axis=0)

    for i, m_val, s_endo in self.regimes():
        m = np.repeat(m_val[None, :], N, axis=0)
        controls_ub[i] = ub_fun(m, s_endo, p)
    return controls_ub

GArray and QP abstractions for dolo.py.

Following Dolo.jl: - QP: Quasi-Point separating discrete grid locators from continuous economic coordinates. - SliceQP: A batch/slice of state points for an entire continuous sub-grid. - GArray: A discrete field bundling a Grid with underlying Data, supporting both geometric indexing and a flat 1D view for numerical solvers.

`GArray` ¶

Represents a discrete field of values living on a Grid.

Supports: - Geometric indexing: garray[loc] accessing a node or slice. - Duality: flat 1D view via ravel() / unravel() for solvers. - Standard array arithmetic and norm.

Source code in src/dolo/numeric/garray.py

class GArray:
    """Represents a discrete field of values living on a Grid.

    Supports:
    - Geometric indexing: garray[loc] accessing a node or slice.
    - Duality: flat 1D view via ravel() / unravel() for solvers.
    - Standard array arithmetic and norm.
    """

    def __init__(self, grid: Any, data: np.ndarray):
        self.grid = grid
        self.data = np.asarray(data)

    @classmethod
    def zeros(cls, grid: Any, n_vars: int, n_regimes: Optional[int] = None) -> GArray:
        """Create a zero GArray conforming to grid structure."""
        if isinstance(grid, ProductGrid) and not isinstance(grid.grids[0], EmptyGrid):
            n_m = grid.grids[0].n_nodes
            n_s = grid.grids[1].n_nodes
            data = np.zeros((n_m, n_s, n_vars))
        elif isinstance(grid, ProductGrid) and isinstance(grid.grids[0], EmptyGrid):
            n_s = grid.grids[1].n_nodes
            data = np.zeros((1, n_s, n_vars))
        elif hasattr(grid, "n_nodes"):
            n_s = grid.n_nodes
            data = np.zeros((1, n_s, n_vars))
        else:
            n_reg = n_regimes if n_regimes is not None else 1
            data = np.zeros((n_reg, len(grid), n_vars))
        return cls(grid, data)

    @classmethod
    def zeros_like(cls, other: GArray) -> GArray:
        """Create a zero GArray matching the grid and shape of another GArray."""
        return cls(other.grid, np.zeros_like(other.data))

    @classmethod
    def ones_like(cls, other: GArray) -> GArray:
        """Create a ones GArray matching the grid and shape of another GArray."""
        return cls(other.grid, np.ones_like(other.data))

    @property
    def shape(self) -> Tuple[int, ...]:
        return self.data.shape

    @property
    def dtype(self) -> np.dtype:
        return self.data.dtype

    def __getitem__(self, loc: Any) -> np.ndarray:
        if isinstance(loc, tuple):
            if len(loc) == 1:
                return self.data[loc[0]]
            return self.data[loc]
        return self.data[loc]

    def __setitem__(self, loc: Any, value: Any) -> None:
        if isinstance(loc, tuple):
            if len(loc) == 1:
                self.data[loc[0]] = value
            else:
                self.data[loc] = value
        else:
            self.data[loc] = value

    def ravel(self) -> np.ndarray:
        """Flat 1D view of the degrees of freedom for numerical solvers."""
        return self.data.ravel()

    def unravel(self, flat_data: np.ndarray) -> GArray:
        """Wrap flat solver iterate back into a structured GArray."""
        return GArray(self.grid, np.asarray(flat_data).reshape(self.data.shape))

    def norm(self) -> float:
        """Maximum absolute norm (L_infinity) across all nodes and variables."""
        return float(np.max(np.abs(self.data)))

    def copy(self) -> GArray:
        return GArray(self.grid, self.data.copy())

    def __sub__(self, other: Union[GArray, np.ndarray, float]) -> GArray:
        other_data = other.data if isinstance(other, GArray) else other
        return GArray(self.grid, self.data - other_data)

    def __add__(self, other: Union[GArray, np.ndarray, float]) -> GArray:
        other_data = other.data if isinstance(other, GArray) else other
        return GArray(self.grid, self.data + other_data)

    def __mul__(self, other: Union[GArray, np.ndarray, float]) -> GArray:
        other_data = other.data if isinstance(other, GArray) else other
        return GArray(self.grid, self.data * other_data)

    def __rmul__(self, scalar: float) -> GArray:
        return GArray(self.grid, self.data * scalar)

    def __truediv__(self, other: Union[GArray, np.ndarray, float]) -> GArray:
        other_data = other.data if isinstance(other, GArray) else other
        return GArray(self.grid, self.data / other_data)

    def __neg__(self) -> GArray:
        return GArray(self.grid, -self.data)

    def __repr__(self) -> str:
        grid_name = type(self.grid).__name__
        return f"GArray(grid={grid_name}, shape={self.data.shape})"

`norm()` ¶

Maximum absolute norm (L_infinity) across all nodes and variables.

Source code in src/dolo/numeric/garray.py

def norm(self) -> float:
    """Maximum absolute norm (L_infinity) across all nodes and variables."""
    return float(np.max(np.abs(self.data)))

`ones_like(other)` `classmethod` ¶

Create a ones GArray matching the grid and shape of another GArray.

Source code in src/dolo/numeric/garray.py

@classmethod
def ones_like(cls, other: GArray) -> GArray:
    """Create a ones GArray matching the grid and shape of another GArray."""
    return cls(other.grid, np.ones_like(other.data))

`ravel()` ¶

Flat 1D view of the degrees of freedom for numerical solvers.

Source code in src/dolo/numeric/garray.py

def ravel(self) -> np.ndarray:
    """Flat 1D view of the degrees of freedom for numerical solvers."""
    return self.data.ravel()

`unravel(flat_data)` ¶

Wrap flat solver iterate back into a structured GArray.

Source code in src/dolo/numeric/garray.py

def unravel(self, flat_data: np.ndarray) -> GArray:
    """Wrap flat solver iterate back into a structured GArray."""
    return GArray(self.grid, np.asarray(flat_data).reshape(self.data.shape))

`zeros(grid, n_vars, n_regimes=None)` `classmethod` ¶

Create a zero GArray conforming to grid structure.

Source code in src/dolo/numeric/garray.py

@classmethod
def zeros(cls, grid: Any, n_vars: int, n_regimes: Optional[int] = None) -> GArray:
    """Create a zero GArray conforming to grid structure."""
    if isinstance(grid, ProductGrid) and not isinstance(grid.grids[0], EmptyGrid):
        n_m = grid.grids[0].n_nodes
        n_s = grid.grids[1].n_nodes
        data = np.zeros((n_m, n_s, n_vars))
    elif isinstance(grid, ProductGrid) and isinstance(grid.grids[0], EmptyGrid):
        n_s = grid.grids[1].n_nodes
        data = np.zeros((1, n_s, n_vars))
    elif hasattr(grid, "n_nodes"):
        n_s = grid.n_nodes
        data = np.zeros((1, n_s, n_vars))
    else:
        n_reg = n_regimes if n_regimes is not None else 1
        data = np.zeros((n_reg, len(grid), n_vars))
    return cls(grid, data)

`zeros_like(other)` `classmethod` ¶

Create a zero GArray matching the grid and shape of another GArray.

Source code in src/dolo/numeric/garray.py

@classmethod
def zeros_like(cls, other: GArray) -> GArray:
    """Create a zero GArray matching the grid and shape of another GArray."""
    return cls(other.grid, np.zeros_like(other.data))

`QualifiedBatch` `dataclass` ¶

A batch/slice of qualified state points for an entire continuous sub-grid.

Attributes¶

loc : Any Target discrete regime index / locator (or None for pure continuous models). exo : np.ndarray Exogenous state values for this branch, shape (N, n_m) or (n_m,). endo : np.ndarray Endogenous state matrix, shape (N, n_s).

Source code in src/dolo/numeric/garray.py

@dataclass
class QualifiedBatch:
    """A batch/slice of qualified state points for an entire continuous sub-grid.

    Attributes
    ----------
    loc : Any
        Target discrete regime index / locator (or None for pure continuous models).
    exo : np.ndarray
        Exogenous state values for this branch, shape (N, n_m) or (n_m,).
    endo : np.ndarray
        Endogenous state matrix, shape (N, n_s).
    """

    loc: Any
    exo: np.ndarray
    endo: np.ndarray

    @property
    def val(self) -> np.ndarray:
        """Full economic state matrix [exo, endo] of shape (N, n_m + n_s)."""
        if self.exo.ndim == 1:
            exo_tiled = np.tile(self.exo, (self.endo.shape[0], 1))
            return np.hstack([exo_tiled, self.endo])
        return np.hstack([self.exo, self.endo])

    @property
    def coords(self) -> np.ndarray:
        """Full economic coordinates matrix."""
        return self.val

`coords` `property` ¶

Full economic coordinates matrix.

`val` `property` ¶

Full economic state matrix [exo, endo] of shape (N, n_m + n_s).

`QualifiedPoint` `dataclass` ¶

A point in a mathematical space qualified with locator context and evaluated values.

A QualifiedPoint bundles: - loc (Locator / Address): How the point is addressed or indexed in the computational domain. This can be discrete (e.g. a regime index i or grid coordinate (i, j) for table lookups) or continuous / differentiable (e.g. normalized coordinates for spline evaluation). - val (Values / Payload): The evaluated economic coordinates and derived quantities (e.g. state vector [z, k], coordinate products, or precomputed basis features) passed to economic equations (arbitrage, transition, utility).

Attributes¶

loc : Any Computational locator / index. val : np.ndarray Evaluated economic coordinates / values.

Source code in src/dolo/numeric/garray.py

@dataclass
class QualifiedPoint:
    """A point in a mathematical space qualified with locator context and evaluated values.

    A QualifiedPoint bundles:
    - `loc` (Locator / Address): How the point is addressed or indexed in the
      computational domain. This can be discrete (e.g. a regime index `i` or grid
      coordinate `(i, j)` for table lookups) or continuous / differentiable
      (e.g. normalized coordinates for spline evaluation).
    - `val` (Values / Payload): The evaluated economic coordinates and derived quantities
      (e.g. state vector `[z, k]`, coordinate products, or precomputed basis features)
      passed to economic equations (arbitrage, transition, utility).

    Attributes
    ----------
    loc : Any
        Computational locator / index.
    val : np.ndarray
        Evaluated economic coordinates / values.
    """

    loc: Any
    val: np.ndarray

    @property
    def coords(self) -> np.ndarray:
        """Alias for economic coordinates `val`."""
        return self.val

    @property
    def index(self) -> Any:
        """Discrete index, if `loc` is a tuple or scalar index."""
        if isinstance(self.loc, tuple) and len(self.loc) > 0:
            return self.loc[0]
        return self.loc

    def __repr__(self) -> str:
        return f"QualifiedPoint(loc={self.loc}, val={np.round(self.val, 4)})"

`coords` `property` ¶

Alias for economic coordinates val.

`index` `property` ¶

Discrete index, if loc is a tuple or scalar index.

Tau transition operator and trembling hand grid projection.

Following Dolo.jl: - Tau: Generates the conditional transition distribution {(w_k, S_k)} over next states. - trembling_hand: Multilinear grid projector mapping continuous next states onto grid nodes. - tau_fit: Transition distribution projected onto discrete grid vertices.

`Tau` ¶

The Tau transition operator.

Encapsulates the model's exogenous stochastic process and endogenous law of motion into a single conditional distribution operator: (s, x) -> {(w_k, S_k)}

Source code in src/dolo/numeric/tau.py

class Tau:
    """The Tau transition operator.

    Encapsulates the model's exogenous stochastic process and endogenous
    law of motion into a single conditional distribution operator:
        (s, x) -> {(w_k, S_k)}
    """

    def __init__(self, g_fun: Any, dprocess: Any, parms: np.ndarray):
        self.g = g_fun
        self.dprocess = dprocess
        self.parms = parms

    def slice_transitions(
        self, i_m: int, m_val: np.ndarray, s_endo: np.ndarray, xm: np.ndarray
    ) -> Iterator[Tuple[float, SliceQP]]:
        """Slice-level transition: vectorized across all N points of continuous grid.

        Yields (prob, SliceQP) where SliceQP contains:
        - loc: target discrete regime index
        - exo: next exogenous shock vector
        - endo: next endogenous state matrix (N, n_s)
        """
        N = s_endo.shape[0]
        n_inodes = self.dprocess.n_inodes(i_m)
        m_tiled = np.tile(m_val, (N, 1))

        for I_m in range(n_inodes):
            prob = float(self.dprocess.iweight(i_m, I_m))
            M_val = self.dprocess.inode(i_m, I_m)
            M_tiled = np.tile(M_val, (N, 1))

            S_endo = self.g(m_tiled, s_endo, xm, M_tiled, self.parms)
            target_regime = 0 if self.dprocess.n_nodes == 1 else I_m
            yield prob, SliceQP(loc=target_regime, exo=M_val, endo=S_endo)

    def point_transitions(self, s: QP, x: np.ndarray) -> Iterator[Tuple[float, QP]]:
        """Point-level transition (pure Julia-style).

        Given a single state point QP(loc, val) and control vector x,
        yields (prob, QP(loc_next, val_next)).
        """
        if isinstance(s.loc, tuple):
            i_m, s_endo = s.loc[0], s.loc[1]
        else:
            i_m, s_endo = 0, s.val

        m_val = self.dprocess.node(i_m)
        n_inodes = self.dprocess.n_inodes(i_m)

        m_1d = m_val[None, :]
        s_1d = s_endo[None, :]
        x_1d = x[None, :]

        for I_m in range(n_inodes):
            prob = float(self.dprocess.iweight(i_m, I_m))
            M_val = self.dprocess.inode(i_m, I_m)
            M_1d = M_val[None, :]

            S_endo = self.g(m_1d, s_1d, x_1d, M_1d, self.parms)[0]
            val_next = np.concatenate([M_val, S_endo])
            target_regime = 0 if self.dprocess.n_nodes == 1 else I_m
            yield prob, QP(loc=(target_regime, S_endo), val=val_next)

    def tau_fit(
        self,
        i_m: int,
        m_val: np.ndarray,
        s_endo: np.ndarray,
        xm: np.ndarray,
        endo_grid: UniformCartesianGrid,
    ) -> Iterator[Tuple[int, float, np.ndarray, np.ndarray]]:
        """Projected transitions on the grid vertices via trembling hand.

        For each shock branch and each multilinear vertex corner, yields:
        (target_regime_j, shock_prob, corner_weights, corner_grid_indices)
        where:
        - corner_weights : np.ndarray of shape (N,)
        - corner_grid_indices : np.ndarray of shape (N,) on endo_grid
        """
        for prob, S in self.slice_transitions(i_m, m_val, s_endo, xm):
            target_regime = S.loc
            projections = trembling_hand(endo_grid, S.endo)
            for w_corner, idx_corner in projections:
                yield target_regime, prob, w_corner, idx_corner

`point_transitions(s, x)` ¶

Point-level transition (pure Julia-style).

Given a single state point QP(loc, val) and control vector x, yields (prob, QP(loc_next, val_next)).

Source code in src/dolo/numeric/tau.py

def point_transitions(self, s: QP, x: np.ndarray) -> Iterator[Tuple[float, QP]]:
    """Point-level transition (pure Julia-style).

    Given a single state point QP(loc, val) and control vector x,
    yields (prob, QP(loc_next, val_next)).
    """
    if isinstance(s.loc, tuple):
        i_m, s_endo = s.loc[0], s.loc[1]
    else:
        i_m, s_endo = 0, s.val

    m_val = self.dprocess.node(i_m)
    n_inodes = self.dprocess.n_inodes(i_m)

    m_1d = m_val[None, :]
    s_1d = s_endo[None, :]
    x_1d = x[None, :]

    for I_m in range(n_inodes):
        prob = float(self.dprocess.iweight(i_m, I_m))
        M_val = self.dprocess.inode(i_m, I_m)
        M_1d = M_val[None, :]

        S_endo = self.g(m_1d, s_1d, x_1d, M_1d, self.parms)[0]
        val_next = np.concatenate([M_val, S_endo])
        target_regime = 0 if self.dprocess.n_nodes == 1 else I_m
        yield prob, QP(loc=(target_regime, S_endo), val=val_next)

`slice_transitions(i_m, m_val, s_endo, xm)` ¶

Slice-level transition: vectorized across all N points of continuous grid.

Yields (prob, SliceQP) where SliceQP contains: - loc: target discrete regime index - exo: next exogenous shock vector - endo: next endogenous state matrix (N, n_s)

Source code in src/dolo/numeric/tau.py

def slice_transitions(
    self, i_m: int, m_val: np.ndarray, s_endo: np.ndarray, xm: np.ndarray
) -> Iterator[Tuple[float, SliceQP]]:
    """Slice-level transition: vectorized across all N points of continuous grid.

    Yields (prob, SliceQP) where SliceQP contains:
    - loc: target discrete regime index
    - exo: next exogenous shock vector
    - endo: next endogenous state matrix (N, n_s)
    """
    N = s_endo.shape[0]
    n_inodes = self.dprocess.n_inodes(i_m)
    m_tiled = np.tile(m_val, (N, 1))

    for I_m in range(n_inodes):
        prob = float(self.dprocess.iweight(i_m, I_m))
        M_val = self.dprocess.inode(i_m, I_m)
        M_tiled = np.tile(M_val, (N, 1))

        S_endo = self.g(m_tiled, s_endo, xm, M_tiled, self.parms)
        target_regime = 0 if self.dprocess.n_nodes == 1 else I_m
        yield prob, SliceQP(loc=target_regime, exo=M_val, endo=S_endo)

`tau_fit(i_m, m_val, s_endo, xm, endo_grid)` ¶

Projected transitions on the grid vertices via trembling hand.

For each shock branch and each multilinear vertex corner, yields: (target_regime_j, shock_prob, corner_weights, corner_grid_indices) where: - corner_weights : np.ndarray of shape (N,) - corner_grid_indices : np.ndarray of shape (N,) on endo_grid

Source code in src/dolo/numeric/tau.py

def tau_fit(
    self,
    i_m: int,
    m_val: np.ndarray,
    s_endo: np.ndarray,
    xm: np.ndarray,
    endo_grid: UniformCartesianGrid,
) -> Iterator[Tuple[int, float, np.ndarray, np.ndarray]]:
    """Projected transitions on the grid vertices via trembling hand.

    For each shock branch and each multilinear vertex corner, yields:
    (target_regime_j, shock_prob, corner_weights, corner_grid_indices)
    where:
    - corner_weights : np.ndarray of shape (N,)
    - corner_grid_indices : np.ndarray of shape (N,) on endo_grid
    """
    for prob, S in self.slice_transitions(i_m, m_val, s_endo, xm):
        target_regime = S.loc
        projections = trembling_hand(endo_grid, S.endo)
        for w_corner, idx_corner in projections:
            yield target_regime, prob, w_corner, idx_corner

`euler_residuals_tau(dmodel, controls, dr)` ¶

Universal Euler residuals written entirely in terms of GArray and Tau.

Source code in src/dolo/numeric/tau.py

def euler_residuals_tau(
    dmodel: Any, controls: GArray, dr: Any
) -> GArray:
    """Universal Euler residuals written entirely in terms of GArray and Tau."""
    res = GArray.zeros_like(controls)
    for i_m, m_val, s_endo in dmodel.regimes():
        xm = controls[i_m]
        for prob, S in dmodel.tau.slice_transitions(i_m, m_val, s_endo, xm):
            XM = dr(S)
            rr = dmodel.arbitrage(m_val, s_endo, xm, S.exo, S.endo, XM)
            res[i_m] += prob * rr
    return res

`trembling_hand(grid, points)` ¶

Multilinear projection of continuous points onto Cartesian grid vertices.

Parameters¶

grid : UniformCartesianGrid The target Cartesian grid with attributes min, max, n. points : np.ndarray, shape (N, d) Continuous points to project onto the grid.

Returns¶

List of (weights, indices) tuples for each of the 2^d vertices: - weights : np.ndarray, shape (N,) - indices : np.ndarray, shape (N,) - linear indices on the Cartesian grid.

Source code in src/dolo/numeric/tau.py

def trembling_hand(
    grid: UniformCartesianGrid, points: np.ndarray
) -> List[Tuple[np.ndarray, np.ndarray]]:
    """Multilinear projection of continuous points onto Cartesian grid vertices.

    Parameters
    ----------
    grid : UniformCartesianGrid
        The target Cartesian grid with attributes `min`, `max`, `n`.
    points : np.ndarray, shape (N, d)
        Continuous points to project onto the grid.

    Returns
    -------
    List of (weights, indices) tuples for each of the 2^d vertices:
    - weights : np.ndarray, shape (N,)
    - indices : np.ndarray, shape (N,) - linear indices on the Cartesian grid.
    """
    d = grid.d
    N = points.shape[0]

    dim_decomps = []
    for k in range(d):
        min_k = float(grid.min[k])
        max_k = float(grid.max[k])
        n_k = int(grid.n[k])
        i0, i1, w0, w1 = trembling_hand_1d(min_k, max_k, n_k, points[:, k])
        dim_decomps.append([(i0, w0), (i1, w1)])

    # Compute striding for linear indices
    strides = [1]
    for k in range(d - 1, 0, -1):
        strides.insert(0, strides[0] * int(grid.n[k]))

    results = []
    # 2^d vertex combinations
    for corner in itertools.product(*dim_decomps):
        # corner is a tuple of (idx_k, weight_k)
        corner_weight = np.ones(N, dtype=float)
        corner_idx = np.zeros(N, dtype=np.int64)
        for k, (idx_k, w_k) in enumerate(corner):
            corner_weight *= w_k
            corner_idx += idx_k * strides[k]
        results.append((corner_weight, corner_idx))

    return results

`trembling_hand_1d(min_val, max_val, n_points, points_1d)` ¶

1D trembling-hand multilinear decomposition.

For an array of values points_1d, finds the two enclosing grid vertices i0, i1 and their linear weights w0 = 1 - lam, w1 = lam.

Returns¶

i0 : np.ndarray (int) Left vertex index (0 to n_points - 2). i1 : np.ndarray (int) Right vertex index (1 to n_points - 1). w0 : np.ndarray (float) Weight for left vertex. w1 : np.ndarray (float) Weight for right vertex.

Source code in src/dolo/numeric/tau.py

def trembling_hand_1d(
    min_val: float, max_val: float, n_points: int, points_1d: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """1D trembling-hand multilinear decomposition.

    For an array of values `points_1d`, finds the two enclosing grid vertices
    `i0`, `i1` and their linear weights `w0 = 1 - lam`, `w1 = lam`.

    Returns
    -------
    i0 : np.ndarray (int)
        Left vertex index (0 to n_points - 2).
    i1 : np.ndarray (int)
        Right vertex index (1 to n_points - 1).
    w0 : np.ndarray (float)
        Weight for left vertex.
    w1 : np.ndarray (float)
        Weight for right vertex.
    """
    scale = (points_1d - min_val) / (max_val - min_val) * (n_points - 1)
    i0 = np.floor(scale).astype(np.int64)
    i0 = np.clip(i0, 0, n_points - 2)
    i1 = i0 + 1
    lam = np.clip(scale - i0, 0.0, 1.0)
    w0 = 1.0 - lam
    w1 = lam
    return i0, i1, w0, w1

Space hierarchy for dolo.py.

Following Dolo.jl: - Space: Mathematical domain for economic variables (continuous or discrete). - CartesianSpace: Continuous bounded hypercube [min, max] (loc = val). - DiscreteSpace (GridSpace): Finite set of discrete points/regimes (loc = index). - ProductSpace: Cartesian product of multiple spaces (e.g. exo × endo). - EmptySpace: Neutral empty domain (discretizes to EmptyGrid).

`CartesianSpace` ¶

Bases: Space, dict

Continuous hypercube domain bounded by [min, max].

Supports dict-like access space['k'] for backwards compatibility with legacy CartesianDomain.

Source code in src/dolo/numeric/spaces.py

class CartesianSpace(Space, dict):
    """Continuous hypercube domain bounded by [min, max].

    Supports dict-like access `space['k']` for backwards compatibility with
    legacy `CartesianDomain`.
    """

    def __init__(
        self,
        bounds: Optional[Dict[str, Sequence[float]]] = None,
        *,
        min: Optional[Sequence[float]] = None,
        max: Optional[Sequence[float]] = None,
        variables: Optional[Sequence[str]] = None,
        **kwargs,
    ):
        dict.__init__(self)

        # Merge dict kwargs or bounds dict
        all_bounds: Dict[str, np.ndarray] = {}
        if bounds is not None:
            for k, v in bounds.items():
                all_bounds[k] = np.array(v, dtype=float)
        for k, v in kwargs.items():
            all_bounds[k] = np.array(v, dtype=float)

        if all_bounds:
            self.variables = tuple(all_bounds.keys())
            self._min = np.array(
                [all_bounds[k][0] for k in self.variables], dtype=float
            )
            self._max = np.array(
                [all_bounds[k][1] for k in self.variables], dtype=float
            )
            for k, v in all_bounds.items():
                self[k] = v
        else:
            if min is None or max is None:
                raise ValueError(
                    "Must specify either bounds dict or min and max arrays."
                )
            self._min = np.array(min, dtype=float)
            self._max = np.array(max, dtype=float)
            if variables is not None:
                self.variables = tuple(variables)
            else:
                self.variables = tuple(f"x_{i}" for i in range(len(self._min)))
            for i, k in enumerate(self.variables):
                self[k] = np.array([self._min[i], self._max[i]], dtype=float)

    @property
    def min(self) -> np.ndarray:
        return self._min

    @property
    def max(self) -> np.ndarray:
        return self._max

    @property
    def states(self) -> Tuple[str, ...]:
        """Backwards compatibility alias for variables."""
        return self.variables

    def contains(self, point: Any) -> bool:
        if isinstance(point, str):
            return point in self.variables
        try:
            pt = np.asarray(point, dtype=float)
            if pt.ndim == 1:
                return bool(np.all(pt >= self._min) and np.all(pt <= self._max))
            return bool(
                np.all(pt >= self._min, axis=-1).all()
                and np.all(pt <= self._max, axis=-1).all()
            )
        except (ValueError, TypeError):
            return False

    def __contains__(self, item: Any) -> bool:
        if isinstance(item, str):
            return dict.__contains__(self, item)
        return self.contains(item)

    def qp(self, loc: Any, val: Optional[Any] = None) -> QualifiedPoint:
        """In CartesianSpace, location and continuous coordinate values coincide."""
        coords = np.asarray(loc if val is None else val, dtype=float)
        return QualifiedPoint(loc=coords, val=coords)

    def draw(self, n: int = 1) -> Union[QualifiedPoint, List[QualifiedPoint]]:
        pts = np.random.uniform(self._min, self._max, size=(n, self.d))
        qps = [self.qp(pts[i]) for i in range(n)]
        return qps[0] if n == 1 else qps

    def discretize(
        self,
        n: Optional[Union[int, Sequence[int]]] = None,
        type: str = "cartesian",
        **kwargs,
    ) -> Grid:
        grid_type = str(type).lower()
        if n is None:
            n = [20] * self.d
        elif isinstance(n, int):
            n = [n] * self.d

        if grid_type in ("cartesian", "uniform", "uniformcartesian"):
            return UniformCartesianGrid(self.min, self.max, n)
        elif grid_type in ("smolyak", "smolyakgrid"):
            mu = kwargs.get("mu", 2)
            return SmolyakGrid(self.min, self.max, mu=mu)
        elif grid_type in ("nonuniform", "nonuniformcartesian"):
            nodes = kwargs.get("nodes")
            if nodes is None:
                raise ValueError("Non-uniform Cartesian grid requires explicit nodes.")
            return NonUniformCartesianGrid(nodes)
        else:
            raise ValueError(f"Unknown grid type: {grid_type}")

    def __repr__(self) -> str:
        b_str = ", ".join(
            f"{k}: [{self[k][0]:.3g}, {self[k][1]:.3g}]" for k in self.variables
        )
        return f"CartesianSpace({b_str})"

`states` `property` ¶

Backwards compatibility alias for variables.

`qp(loc, val=None)` ¶

In CartesianSpace, location and continuous coordinate values coincide.

Source code in src/dolo/numeric/spaces.py

def qp(self, loc: Any, val: Optional[Any] = None) -> QualifiedPoint:
    """In CartesianSpace, location and continuous coordinate values coincide."""
    coords = np.asarray(loc if val is None else val, dtype=float)
    return QualifiedPoint(loc=coords, val=coords)

`DiscreteSpace` ¶

Bases: Space

Finite discrete set of points or regimes (e.g. Markov states).

In a DiscreteSpace: - loc is the integer regime index i (0 <= i < n_points). - val is the discrete coordinate vector points[i].

Source code in src/dolo/numeric/spaces.py

class DiscreteSpace(Space):
    """Finite discrete set of points or regimes (e.g. Markov states).

    In a DiscreteSpace:
    - `loc` is the integer regime index `i` (0 <= i < n_points).
    - `val` is the discrete coordinate vector `points[i]`.
    """

    def __init__(
        self,
        points: Union[Sequence[Sequence[float]], np.ndarray],
        variables: Optional[Sequence[str]] = None,
    ):
        self.points = np.asarray(points, dtype=float)
        if self.points.ndim == 1:
            self.points = self.points[:, None]
        self.n_points = self.points.shape[0]

        if variables is not None:
            self.variables = tuple(variables)
        else:
            self.variables = tuple(f"m_{i}" for i in range(self.points.shape[1]))

    def contains(self, point: Any) -> bool:
        if isinstance(point, (int, np.integer)):
            return 0 <= int(point) < self.n_points
        pt = np.asarray(point, dtype=float)
        diffs = np.linalg.norm(self.points - pt, axis=-1)
        return bool(np.any(diffs < 1e-12))

    def qp(self, loc: Any, val: Optional[Any] = None) -> QualifiedPoint:
        if isinstance(loc, (int, np.integer)):
            idx = int(loc)
            coord = self.points[idx] if val is None else np.asarray(val, dtype=float)
            return QualifiedPoint(loc=idx, val=coord)
        pt = np.asarray(loc, dtype=float)
        diffs = np.linalg.norm(self.points - pt, axis=-1)
        idx = int(np.argmin(diffs))
        return QualifiedPoint(
            loc=idx, val=pt if val is None else np.asarray(val, dtype=float)
        )

    def draw(self, n: int = 1) -> Union[QualifiedPoint, List[QualifiedPoint]]:
        indices = np.random.choice(self.n_points, size=n)
        qps = [self.qp(int(idx)) for idx in indices]
        return qps[0] if n == 1 else qps

    def discretize(self, **options) -> UnstructuredGrid:
        return UnstructuredGrid(self.points)

    def __repr__(self) -> str:
        return f"DiscreteSpace(n_points={self.n_points}, vars={self.variables})"

`EmptySpace` ¶

Bases: Space

Empty / Neutral space (used when exogenous shocks are absent).

Source code in src/dolo/numeric/spaces.py

class EmptySpace(Space):
    """Empty / Neutral space (used when exogenous shocks are absent)."""

    variables: Tuple[str, ...] = ()

    def contains(self, point: Any) -> bool:
        return True

    def qp(self, loc: Any = None, val: Optional[Any] = None) -> QualifiedPoint:
        return QualifiedPoint(loc=None, val=np.empty((0,)))

    def draw(self, n: int = 1) -> Union[QualifiedPoint, List[QualifiedPoint]]:
        qps = [self.qp() for _ in range(n)]
        return qps[0] if n == 1 else qps

    def discretize(self, **options) -> EmptyGrid:
        return EmptyGrid()

    def __repr__(self) -> str:
        return "EmptySpace()"

`ProductSpace` ¶

Bases: Space

Product of two or more spaces (e.g. Exogenous × Endogenous).

Source code in src/dolo/numeric/spaces.py

class ProductSpace(Space):
    """Product of two or more spaces (e.g. Exogenous × Endogenous)."""

    def __init__(
        self,
        spaces: Sequence[Space],
        names: Optional[Sequence[str]] = None,
    ):
        self.spaces = tuple(spaces)
        if names is not None:
            self.names = tuple(names)
        elif len(self.spaces) == 2:
            self.names = ("exo", "endo")
        else:
            self.names = tuple(f"space_{i}" for i in range(len(self.spaces)))

        all_vars: List[str] = []
        for s in self.spaces:
            all_vars.extend(s.variables)
        self.variables = tuple(all_vars)

    @property
    def exo(self) -> Space:
        return self.spaces[0]

    @property
    def endo(self) -> Space:
        return self.spaces[1]

    def __getitem__(self, key: Union[int, str]) -> Space:
        if isinstance(key, str):
            idx = self.names.index(key)
            return self.spaces[idx]
        return self.spaces[key]

    def contains(self, point: Any) -> bool:
        if isinstance(point, (tuple, list)) and len(point) == len(self.spaces):
            return all(s.contains(p) for s, p in zip(self.spaces, point))
        return False

    def qp(self, loc: Any, val: Optional[Any] = None) -> QualifiedPoint:
        if isinstance(loc, tuple) and len(loc) == len(self.spaces):
            loc_tuple = tuple(loc)
            if val is not None:
                val_arr = np.asarray(val, dtype=float)
            else:
                qps = [s.qp(l) for s, l in zip(self.spaces, loc)]
                val_arr = np.concatenate([q.val for q in qps])
            return QualifiedPoint(loc=loc_tuple, val=val_arr)
        val_arr = np.asarray(val if val is not None else loc, dtype=float)
        return QualifiedPoint(loc=loc, val=val_arr)

    def draw(self, n: int = 1) -> Union[QualifiedPoint, List[QualifiedPoint]]:
        res = []
        for _ in range(n):
            sub_qps = [s.draw(1) for s in self.spaces]
            loc = tuple(q.loc for q in sub_qps)
            val = np.concatenate([q.val for q in sub_qps])
            res.append(QualifiedPoint(loc=loc, val=val))
        return res[0] if n == 1 else res

    def discretize(self, **options) -> ProductGrid:
        grids = []
        for name, space in zip(self.names, self.spaces):
            sub_opts = (
                options.get(name, {}) if isinstance(options.get(name), dict) else {}
            )
            grids.append(space.discretize(**sub_opts))
        return ProductGrid(grids[0], grids[1], names=list(self.names))

    def __repr__(self) -> str:
        return " × ".join(repr(s) for s in self.spaces)

`Space` ¶

Bases: ABC

Abstract base class representing a mathematical space for economic variables.

Source code in src/dolo/numeric/spaces.py

class Space(ABC):
    """Abstract base class representing a mathematical space for economic variables."""

    variables: Tuple[str, ...]

    @property
    def d(self) -> int:
        """Dimension of the space."""
        return len(self.variables)

    @abstractmethod
    def contains(self, point: Any) -> bool:
        """Check if a point is within this space."""
        pass

    def __contains__(self, point: Any) -> bool:
        return self.contains(point)

    @abstractmethod
    def qp(self, loc: Any, val: Optional[Any] = None) -> QualifiedPoint:
        """Construct a QualifiedPoint bound to this space."""
        pass

    @abstractmethod
    def draw(self, n: int = 1) -> Union[QualifiedPoint, List[QualifiedPoint]]:
        """Sample random QualifiedPoint(s) from this space."""
        pass

    @abstractmethod
    def discretize(self, **options) -> Grid:
        """Discretize the continuous or discrete space into a computational Grid."""
        pass

    def __mul__(self, other: Space) -> ProductSpace:
        if isinstance(other, ProductSpace):
            return ProductSpace((self, *other.spaces))
        return ProductSpace((self, other))

`d` `property` ¶

Dimension of the space.

`contains(point)` `abstractmethod` ¶

Check if a point is within this space.

Source code in src/dolo/numeric/spaces.py

@abstractmethod
def contains(self, point: Any) -> bool:
    """Check if a point is within this space."""
    pass

`discretize(**options)` `abstractmethod` ¶

Discretize the continuous or discrete space into a computational Grid.

Source code in src/dolo/numeric/spaces.py

@abstractmethod
def discretize(self, **options) -> Grid:
    """Discretize the continuous or discrete space into a computational Grid."""
    pass

`draw(n=1)` `abstractmethod` ¶

Sample random QualifiedPoint(s) from this space.

Source code in src/dolo/numeric/spaces.py

@abstractmethod
def draw(self, n: int = 1) -> Union[QualifiedPoint, List[QualifiedPoint]]:
    """Sample random QualifiedPoint(s) from this space."""
    pass

`qp(loc, val=None)` `abstractmethod` ¶

Construct a QualifiedPoint bound to this space.

Source code in src/dolo/numeric/spaces.py

@abstractmethod
def qp(self, loc: Any, val: Optional[Any] = None) -> QualifiedPoint:
    """Construct a QualifiedPoint bound to this space."""
    pass

Approximator hierarchy for dolo.py.

Following Dolo.jl: - Approximator: Operator knowing how to fit coefficients from values on a Grid, and evaluate approximations at query points (QualifiedPoint, QualifiedBatch, or coords). - LinearApproximator: Piecewise multilinear interpolation on CartesianGrid. - CubicSplineApproximator: Prefiltered cubic spline interpolation on CartesianGrid. - SmolyakApproximator: Multivariate Chebyshev polynomial basis on SmolyakGrid. - ProductApproximator: Regime-wise composite approximator over ProductGrid / ProductSpace.