How to use quantecon - 10 common examples

def _solve_via_vfi(jv):
    "compute policy rules via value function iteration"
    v_init = jv.x_grid * 0.6
    V = compute_fixed_point(jv.bellman_operator, v_init,
                            max_iter=3000,
                            error_tol=1e-5)
    return V

def _solve_via_vfi(cp, v_init, return_both=False):
    "compute policy rule using value function iteration"
    v = compute_fixed_point(cp.bellman_operator, v_init, verbose=False,
                            error_tol=1e-5,
                            max_iter=1000)

    # Run one more time to get the policy
    p = cp.bellman_operator(v, return_policy=True)

    if return_both:
        return v, p
    else:
        return p

"gets a new set of solution objects and updates the data file"

    # compute value function and policy rule using vfi
    v_init = np.zeros(len(sp.grid_points)) + sp.c / (1 - sp.beta)
    v = compute_fixed_point(sp.bellman_operator, v_init, error_tol=_tol,
                            max_iter=5000)
    phi_vfi = sp.get_greedy(v)

    # also run v through bellman so I can test if it is a fixed point
    # bellman_operator takes a long time, so store result instead of compute
    new_v = sp.bellman_operator(v)

    # compute policy rule using pfi

    phi_init = np.ones(len(sp.pi_grid))
    phi_pfi = compute_fixed_point(sp.res_wage_operator, phi_init,
                                  error_tol=_tol, max_iter=5000)

    # write all arrays to file
    write_array(f, grp, v, "v")
    write_array(f, grp, phi_vfi, "phi_vfi")
    write_array(f, grp, phi_pfi, "phi_pfi")
    write_array(f, grp, new_v, "new_v")

    # return data
    return v, phi_vfi, phi_pfi, new_v

def _new_solutions(cp, f, grp, which="both"):
    v_init, c_init = cp.initialize()
    if which == "both":

        v_vfi, c_vfi = _solve_via_vfi(cp, v_init, return_both=True)
        c_pfi = _solve_via_pfi(cp, c_init)

        # Store solutions in chunked arrays...
        write_array(f, grp, c_vfi, "c_vfi")
        write_array(f, grp, v_vfi, "v_vfi")
        write_array(f, grp, c_pfi, "c_pfi")

        return v_vfi, c_vfi, c_pfi

    elif which == "vfi":
        v_vfi, c_vfi = _solve_via_vfi(cp, v_init, return_both=True)
        write_array(f, grp, c_vfi, "c_vfi")
        write_array(f, grp, v_vfi, "v_vfi")

        return v_vfi, c_vfi

    elif which == "pfi":
        c_pfi = _solve_via_pfi(cp, c_init)
        write_array(f, grp, c_pfi, "c_pfi")

phi_vfi = sp.get_greedy(v)

    # also run v through bellman so I can test if it is a fixed point
    # bellman_operator takes a long time, so store result instead of compute
    new_v = sp.bellman_operator(v)

    # compute policy rule using pfi

    phi_init = np.ones(len(sp.pi_grid))
    phi_pfi = compute_fixed_point(sp.res_wage_operator, phi_init,
                                  error_tol=_tol, max_iter=5000)

    # write all arrays to file
    write_array(f, grp, v, "v")
    write_array(f, grp, phi_vfi, "phi_vfi")
    write_array(f, grp, phi_pfi, "phi_pfi")
    write_array(f, grp, new_v, "new_v")

    # return data
    return v, phi_vfi, phi_pfi, new_v

def _new_solution(jv, f, grp):
    "gets new solution and updates data file"
    V = _solve_via_vfi(jv)
    write_array(f, grp, V, v_nm)

    return V

# also run v through bellman so I can test if it is a fixed point
    # bellman_operator takes a long time, so store result instead of compute
    new_v = sp.bellman_operator(v)

    # compute policy rule using pfi

    phi_init = np.ones(len(sp.pi_grid))
    phi_pfi = compute_fixed_point(sp.res_wage_operator, phi_init,
                                  error_tol=_tol, max_iter=5000)

    # write all arrays to file
    write_array(f, grp, v, "v")
    write_array(f, grp, phi_vfi, "phi_vfi")
    write_array(f, grp, phi_pfi, "phi_pfi")
    write_array(f, grp, new_v, "new_v")

    # return data
    return v, phi_vfi, phi_pfi, new_v

def _new_solutions(cp, f, grp, which="both"):
    v_init, c_init = cp.initialize()
    if which == "both":

        v_vfi, c_vfi = _solve_via_vfi(cp, v_init, return_both=True)
        c_pfi = _solve_via_pfi(cp, c_init)

        # Store solutions in chunked arrays...
        write_array(f, grp, c_vfi, "c_vfi")
        write_array(f, grp, v_vfi, "v_vfi")
        write_array(f, grp, c_pfi, "c_pfi")

        return v_vfi, c_vfi, c_pfi

    elif which == "vfi":
        v_vfi, c_vfi = _solve_via_vfi(cp, v_init, return_both=True)
        write_array(f, grp, c_vfi, "c_vfi")
        write_array(f, grp, v_vfi, "v_vfi")

        return v_vfi, c_vfi

    elif which == "pfi":
        c_pfi = _solve_via_pfi(cp, c_init)
        write_array(f, grp, c_pfi, "c_pfi")

        return c_pfi

def _new_solutions(cp, f, grp, which="both"):
    v_init, c_init = cp.initialize()
    if which == "both":

        v_vfi, c_vfi = _solve_via_vfi(cp, v_init, return_both=True)
        c_pfi = _solve_via_pfi(cp, c_init)

        # Store solutions in chunked arrays...
        write_array(f, grp, c_vfi, "c_vfi")
        write_array(f, grp, v_vfi, "v_vfi")
        write_array(f, grp, c_pfi, "c_pfi")

        return v_vfi, c_vfi, c_pfi

    elif which == "vfi":
        v_vfi, c_vfi = _solve_via_vfi(cp, v_init, return_both=True)
        write_array(f, grp, c_vfi, "c_vfi")
        write_array(f, grp, v_vfi, "v_vfi")

        return v_vfi, c_vfi

    elif which == "pfi":
        c_pfi = _solve_via_pfi(cp, c_init)
        write_array(f, grp, c_pfi, "c_pfi")

        return c_pfi

def _get_data(sp, force_new=False):
    "get solution data from file, or create if necessary"
    with get_h5_data_file() as f:
        existed, grp = get_h5_data_group("odu")

        if force_new or not existed:
            if existed:
                grp.v._f_remove()
                grp.phi_vfi._f_remove()
                grp.phi_pfi._f_remove()
                grp.new_v._f_remove()
            v, phi_vfi, phi_pfi, new_v = _new_solution(sp, f, grp)

            return v, phi_vfi, phi_pfi, new_v

        # if we made it here, the group exists and we should try to read
        # existing solutions
        try:
            # Try reading data