import mmf.math.multigrid.notes.multigrid_notes as MN
mg = MN.MultigridNeumann(inclusive=False, prob=1)
x = mg.x(); dx = x[1]-x[0]
f = mg.problem.f(x)
b = mg.problem.b(x)
A = mg.to_mat(mg.A, len(f))
f_ = np.linalg.solve(A, b)

g = 0*x
res = [abs(mg.residue(g, b)).max()]
err = [abs(g - f).max()]
err_ = [abs(g - f_).max()]
for n in xrange(10):
   g = mg.f_cycle(g, b)
   res.append(abs(mg.residue(g, b)).max())
   err.append(abs(g - f).max())
   err_.append(abs(g - f_).max())

plt.figure(figsize=(16,6))
plt.subplot(121)
plt.semilogy(res, ':', label='residue')
plt.semilogy(err_, '-', label='err')
plt.semilogy(err, '--', label='exact err')
plt.legend()
plt.title('Not Inclusive')

mg = MN.MultigridNeumann(inclusive=True, prob=1)
x = mg.x(); dx = x[1]-x[0]
f = mg.problem.f(x)
b = mg.problem.b(x)
A = mg.to_mat(mg.A, len(f))
b[[0,-1]] /= 2                      # Correct the endpoints
f_ = np.linalg.solve(A, b)
g = 0*x
res = [abs(mg.residue(g, b)).max()]
err = [abs(g - f).max()]
err_ = [abs(g - f_).max()]
for n in xrange(10):
   g = mg.f_cycle(g, b)
   res.append(abs(mg.residue(g, b)).max())
   err.append(abs(g - f).max())
   err_.append(abs(g - f_).max())

plt.subplot(122)
plt.semilogy(res, ':', label='residue')
plt.semilogy(err_, '-', label='err')
plt.semilogy(err, '--', label='exact err')
plt.legend()
plt.title('Inclusive')