def _new_solution(sp, f, grp):
    "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 _solve_via_pfi(cp, c_init):
    "compute policy rule using policy function iteration"
    p = compute_fixed_point(cp.coleman_operator, c_init, verbose=False,
                            error_tol=1e-5,
                            max_iter=1000)

    return p

 def test_contraction_1(self):
     "compute_fp: convergence inside interval of convergence"
     f = lambda x: self.T(x, self.mu_1)
     for i in self.unit_inverval:
         # should have fixed point of 0.0
         self.assertTrue(abs(compute_fixed_point(f, i, **self.kwargs))
                         < 1e-4)

 def test_not_contraction_2(self):
     "compute_fp: no convergence outside interval of convergence"
     f = lambda x: self.T(x, self.mu_2)
     for i in self.unit_inverval:
         # This shouldn't converge to 0.0
         self.assertFalse(abs(compute_fixed_point(f, i, **self.kwargs))
                          < 1e-4)

    def compute_lt_price(self, error_tol=1e-3, max_iter=50, verbose=0):
        """
        Compute the equilibrium price function associated with Lucas
        tree lt

        Parameters
        ----------
        error_tol, max_iter, verbose
            Arguments to be passed directly to
            `quantecon.compute_fixed_point`. See that docstring for more
            information

        Returns
        -------
        price : array_like(float)
            The prices at the grid points in the attribute `grid` of the
            object

        """
        # == simplify notation == #
        grid, grid_size = self.grid, self.grid_size
        lucas_operator, gamma = self.lucas_operator, self.gamma

        # == Create storage array for compute_fixed_point. Reduces  memory
        # allocation and speeds code up == #
        Tf = np.empty(grid_size)

        # == Initial guess, just a vector of zeros == #
        f_init = np.zeros(grid_size)
        f = compute_fixed_point(lucas_operator, f_init, error_tol,
                                max_iter, verbose, Tf=Tf)

        price = f * grid**gamma

        return price

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 test_not_contraction_1(self):
     "compute_fp: no convergence outside interval of convergence"
     f = lambda x: self.T(x, self.mu_1)
     fp = (4 * self.mu_1 - 1) / (4 * self.mu_1)
     for i in self.unit_inverval:
         # This should not converge  (b/c unique fp is 0.0)
         self.assertFalse(abs(compute_fixed_point(f, i, **self.kwargs)-fp)
                          < 1e-4)

 def test_contraction_2(self):
     "compute_fp: convergence inside interval of convergence"
     f = lambda x: self.T(x, self.mu_2)
     fp = (4 * self.mu_2 - 1) / (4 * self.mu_2)
     for i in self.unit_inverval:
         # This should converge to fp
         self.assertTrue(abs(compute_fixed_point(f, i, **self.kwargs)-fp)
                         < 1e-4)

    def setUpClass(cls):
        jv = JvWorker(A=A, alpha=alpha, beta=beta, grid_size=grid_size)
        cls.jv = jv

        # compute solution
        v_init = _get_vf_guess(jv)
        cls.V = compute_fixed_point(jv.bellman_operator, v_init)
        cls.s_pol, cls.phi_pol = jv.bellman_operator(cls.V * 0.999,
                                                     return_policies=True)

    def test_2d_input(self):
        error_tol = self.coeff**4

        for method in self.methods:
            init = np.array([[-1, 0.5], [-1/3, 0.1]])
            fp_computed = compute_fixed_point(self.f, init,
                                              error_tol=error_tol,
                                              method=method)
            ok_((fp_computed <= error_tol * 2).all())

    def test_num_iter_one(self):
        init = 1.
        error_tol = self.coeff

        for method in self.methods:
            fp_computed = compute_fixed_point(self.f, init,
                                              error_tol=error_tol,
                                              method=method)
            ok_(fp_computed <= error_tol * 2)

    def test_imitation_game_method(self):
        "compute_fp: Test imitation game method"
        method = 'imitation_game'
        error_tol = self.kwargs['error_tol']

        for mu in [self.mu_1, self.mu_2]:
            for i in self.unit_inverval:
                fp_computed = compute_fixed_point(self.T, i, method=method,
                                                  mu=mu, **self.kwargs)
                self.assertTrue(
                    abs(self.T(fp_computed, mu=mu) - fp_computed) <= error_tol
                )

            # numpy array input
            i = np.asarray(self.unit_inverval)
            fp_computed = compute_fixed_point(self.T, i, method=method, mu=mu,
                                              **self.kwargs)
            self.assertTrue(
                abs(self.T(fp_computed, mu=mu) - fp_computed).max() <=
                error_tol
            )

    def test_num_iter_large(self):
        init = 1.
        buff_size = 2**8  # buff_size in 'imitation_game'
        max_iter = buff_size + 2
        error_tol = self.coeff**max_iter

        for method in self.methods:
            fp_computed = compute_fixed_point(self.f, init,
                                              error_tol=error_tol,
                                              max_iter=max_iter, method=method,
                                              print_skip=max_iter)
            ok_(fp_computed <= error_tol * 2)

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

def _new_solution(gm, f, grp):
    "gets a new set of solution objects and updates the data file"

    # compute value function and policy rule using vfi
    v_init = 5 * gm.u(gm.grid) - 25
    v = compute_fixed_point(gm.bellman_operator, v_init, error_tol=_tol,
                            max_iter=5000)
    # sigma = gm.get_greedy(v)

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

    # return data
    return v

def all_param_interact(c, L0, L1, a0, b0, a1, b1, m):
    f0 = np.clip(st.beta.pdf(np.linspace(0, 1, m), a=a0, b=b0), 1e-6, np.inf)
    f0 = f0 / np.sum(f0)
    f1 = np.clip(st.beta.pdf(np.linspace(0, 1, m), a=a1, b=b1), 1e-6, np.inf)
    f1 = f1 / np.sum(f1)  # Make sure sums to 1

    # Create an instance of our WaldFriedman class
    wf = WaldFriedman(c, L0, L1, f0, f1, m=m)
    # Solve using qe's `compute_fixed_point` function
    J = qe.compute_fixed_point(wf.bellman_operator, np.zeros(m),
                               error_tol=1e-7, verbose=False,
                               print_skip=10, max_iter=500)
    lb, ub = wf.find_cutoff_rule(J)

    # Get draws
    ndraws = 500
    cdist, tdist = wf.stopping_dist(ndraws=ndraws)

    fig, ax = plt.subplots(2, 2, figsize=(22, 14))

    ax[0, 0].plot(f0, marker="o", markersize=2.5, linestyle="None", label=r"$f_0$")
    ax[0, 0].plot(f1, marker="o", markersize=2.5, linestyle="None", label=r"$f_1$")
    ax[0, 0].set_ylabel(r"Probability of $z_k$")
    ax[0, 0].set_xlabel(r"$k$")
    ax[0, 0].set_title("Distributions over Outcomes", size=24)

    ax[0, 1].plot(wf.pgrid, J)
    ax[0, 1].annotate(r"$\alpha$", xy=(lb+0.025, 0.5), size=14)
    ax[0, 1].annotate(r"$\beta$", xy=(ub+0.025, 0.5), size=14)
    ax[0, 1].vlines(lb, 0.0, wf.payoff_choose_f1(lb), linestyle="--")
    ax[0, 1].vlines(ub, 0.0, wf.payoff_choose_f0(ub), linestyle="--")
    ax[0, 1].set_ylim(0, 0.5*max(L0, L1))
    ax[0, 1].set_ylabel("Value of Bellman")
    ax[0, 1].set_xlabel(r"$p_k$")
    ax[0, 1].set_title("Bellman Equation", size=24)

    ax[1, 0].hist(tdist, bins=np.max(tdist))
    ax[1, 0].set_title("Stopping Times", size=24)
    ax[1, 0].set_xlabel("Time")
    ax[1, 0].set_ylabel("Density")

    ax[1, 1].hist(cdist, bins=2)
    ax[1, 1].set_title("Correct Decisions", size=24)
    ax[1, 1].annotate("Percent Correct p={}".format(np.mean(cdist)), xy=(0.05, ndraws/2), size=18)

    fig.tight_layout()
    fig.show()

def WaldFriedman_Interactive(m):
    # NOTE: Could add sliders over other variables
    #       as well, but only doing over n for now

    # Choose parameters
    c = 1.25
    L0 = 25.0
    L1 = 25.0

    # Choose n points and distributions
    f0 = np.clip(st.beta.pdf(np.linspace(0, 1, m), a=2.0, b=2.5), 1e-6, np.inf)
    f0 = f0 / np.sum(f0)
    f1 = np.clip(st.beta.pdf(np.linspace(0, 1, m), a=2.5, b=2.0), 1e-6, np.inf)
    f1 = f1 / np.sum(f1)  # Make sure sums to 1

    # Create WaldFriedman class
    wf = WaldFriedman(c, L0, L1, f0, f1, m=m)

    # Solve via VFI
    # Solve using qe's `compute_fixed_point` function
    J = qe.compute_fixed_point(wf.bellman_operator, np.zeros(m),
                               error_tol=1e-7, verbose=False,
                               max_iter=1000)

    lb, ub = wf.find_cutoff_rule(J)

    # Plot
    fig, ax = plt.subplots(figsize=(8, 6))

    fig.suptitle("Value function", size=18)
    ax.set_xlabel("Probability of Model 0")
    ax.set_ylabel("Value Function")

    ax.set_xlim(0, 1.0)
    ax.set_ylim(0, 0.5 * max(L0, L1))
    ax.plot(wf.pgrid, J)

    ax.annotate(r"$\beta$", xy=(ub+0.025, 0.5), size=14)
    ax.annotate(r"$\alpha$", xy=(lb+0.025, 0.5), size=14)
    ax.vlines(lb, 0.0, wf.payoff_choose_f1(lb), linestyle="--")
    ax.vlines(ub, 0.0, wf.payoff_choose_f0(ub), linestyle="--")

    fig.show()

def compute_value_function_cached(grid, beta, alpha, shocks):
    """
    Compute the value function by iterating on the Bellman operator.
    The work is done by QuantEcon's compute_fixed_point function.
    """
    Tw = np.empty(len(grid))
    initial_w = 5 * np.log(grid) - 25

    v_star = compute_fixed_point(bellman_operator, 
            initial_w, 
            1e-4,  # error_tol
            100,   # max_iter
            True,  # verbose
            5,     # print_skip
            grid,
            beta,
            np.log,
            lambda k: k**alpha,
            shocks,
            Tw=Tw,
            compute_policy=False)
    return v_star

L0 = 25
L1 = 25
a0, b0 = 2.5, 2.0
a1, b1 = 2.0, 2.5
m = 25

f0 = np.clip(st.beta.pdf(np.linspace(0, 1, m), a=a0, b=b0), 1e-6, np.inf)
f0 = f0 / np.sum(f0)
f1 = np.clip(st.beta.pdf(np.linspace(0, 1, m), a=a1, b=b1), 1e-6, np.inf)
f1 = f1 / np.sum(f1)  # Make sure sums to 1

# Create an instance of our WaldFriedman class
wf = WaldFriedman(c, L0, L1, f0, f1, m=m)
# Solve using qe's `compute_fixed_point` function
J = qe.compute_fixed_point(wf.bellman_operator, np.zeros(m),
                           error_tol=1e-7, verbose=False,
                           print_skip=10, max_iter=500)
lb, ub = wf.find_cutoff_rule(J)

# Get draws
ndraws = 500
cdist, tdist = wf.stopping_dist(ndraws=ndraws)

fig, ax = plt.subplots(2, 2, figsize=(12, 9))

ax[0, 0].plot(f0, label=r"$f_0$")
ax[0, 0].plot(f1, label=r"$f_1$")
ax[0, 0].set_ylabel(r"probability of $z_k$", size=14)
ax[0, 0].set_xlabel(r"$k$", size=14)
ax[0, 0].set_title("Distributions", size=14)
ax[0, 0].legend(fontsize=14)

def test_raises_value_error_nonpositive_max_iter():
    f = lambda x: 0.5*x
    init = 1.
    max_iter = 0
    fp = compute_fixed_point(f, init, max_iter=max_iter)