Inspecting the solution¶

def tabulate(
    model, dr, state, bounds=None, n_steps=100, s0=None, i0=None, m0=None, **kwargs
):

    import numpy

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

    states_names = model.symbols["states"]
    controls_names = model.symbols["controls"]
    index = states_names.index(str(state))

    if bounds is None:
        try:
            endo_grid = dr.endo_grid
            bounds = [endo_grid.min[index], endo_grid.max[index]]
        except:
            domain = model.domain
            bounds = [domain.min[index], domain.max[index]]
        if bounds is None:
            raise Exception("No bounds provided for simulation or by model.")

    values = numpy.linspace(bounds[0], bounds[1], n_steps)

    if s0 is None:
        s0 = model.calibration["states"]

    svec = numpy.vstack([s0] * n_steps)
    svec[:, index] = values

    try:
        dp = dr.dprocess
    except:
        dp = model.exogenous.discretize()

    if (i0 is None) and (m0 is None):
        from dolo.numeric.grids import UnstructuredGrid

        if isinstance(dp.grid, UnstructuredGrid):
            n_ms = dp.n_nodes
            [q, r] = divmod(n_ms, 2)
            i0 = q - 1 + r
        else:
            m0 = model.calibration["exogenous"]

    if i0 is not None:
        m = dp.node(i0)
        xvec = dr(i0, svec)
    elif m0 is not None:
        m = m0
        xvec = dr(m0, svec)

    mm = numpy.vstack([m] * n_steps)
    l = [mm, svec, xvec]

    series = (
        model.symbols["exogenous"] + model.symbols["states"] + model.symbols["controls"]
    )

    if "auxiliary" in model.functions:
        p = model.calibration["parameters"]
        pp = numpy.vstack([p] * n_steps)
        avec = model.functions["auxiliary"](mm, svec, xvec, pp)
        l.append(avec)
        series.extend(model.symbols["auxiliaries"])

    import pandas

    tb = numpy.concatenate(l, axis=1)
    df = pandas.DataFrame(tb, columns=series)

    return df

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

def response(model, dr, varname, T=40, impulse: float = None):

    i_exo = model.symbols["exogenous"].index(varname)

    if impulse is None:
        try:
            impulse = numpy.sqrt(
                model.exogenous.Σ[i_exo, i_exo]
            )  # works only for IID/AR1
        except:
            impulse = numpy.sqrt(model.exogenous.σ)  # works only for IID/AR1

    e1 = numpy.zeros(len(model.symbols["exogenous"]))
    e1[i_exo] = impulse

    exogenous = model.exogenous
    print(exogenous)
    print(T, e1)
    m_simul = model.exogenous.response(T - 1, e1)  # this is an xarray T x V
    m_simul = m_simul.expand_dims("N")
    m_simul = m_simul.transpose("T", "N", "V").data

    sim = simulate(model, dr, N=1, T=T, driving_process=m_simul, stochastic=False)

    irf = sim.sel(N=0)

    return irf

def plot_irfs(
    sims,
    variables=None,
    titles=None,
    layout=None,
    horizon=None,
    figsize=None,
    plot_options={},
    line_options=None,
):

    from matplotlib import pyplot as plt

    if not isinstance(sims, list):
        # we have only one series to compare
        sims = [sims]

    if variables is None:
        variables = sims[0].columns

    if horizon is None:
        horizon = sims[0].shape[0]

    if titles is None:
        titles = variables

    n_v = len(variables)
    max_v = max(default_layout.keys())
    if n_v > max_v:
        raise Exception("Too many graphs")

    if layout is None:
        if figsize is None:
            figsize, layout = default_layout[n_v]
        else:
            layout = default_layout[n_v][1]

    if line_options is None:
        line_options = [{}] * n_v

    fig, axes = plt.subplots(layout[0], layout[1], figsize=figsize)
    axes = axes.ravel()
    for i, v in enumerate(variables):
        ax = axes[i]
        for j, sim in enumerate(sims):
            arguments = {}
            arguments.update(plot_options)
            arguments.update(line_options[j])
            ax.plot(sim[v][:horizon], **arguments)
        ax.set_title(titles[i])
        ax.set_xlim(0, horizon)
        max_ = max([sim[v].max() for sim in sims])
        min_ = min([sim[v].min() for sim in sims])
        if (max_ - min_) > 0.001:  # pretty arbitrary
            margin = (max_ - min_) * 0.05  # 5% on each side
            ax.set_ylim(min_ - margin, max_ + margin)
        ax.grid(True)
    # fig.tight_layout()

    return fig