import scipy as sp
import numpy as np
import scipy.sparse as sparse
import scipy.sparse.linalg as sla
import matplotlib.pyplot as plt
%matplotlib inline

def hnorm(r):
    """define ||r||_h = h ||r||_2"""
    n = len(r)
    h = 1.0 / (n+1)
    hrnorm = h * np.linalg.norm(r)
    return hrnorm

def poissonop(n):
    """
    Poisson operator h^{-2} * [-1 2 1]
    """
    A = (n+1)**2 * sparse.diags([-1, 2, -1], [-1, 0, 1], shape=(n,n), format='csr')
    return A

def acoeff(x):
    k = 50
    rho = 0.95
    return 1 + rho*np.sin(k*np.pi*x)

def poissonop2(n):
    xx = np.linspace(0,1, n+2)
    xxhalf = (xx[:-1] + xx[1:]) / 2
    ahalf = acoeff(xxhalf)
    diagonals = [-ahalf[1:-1], ahalf[:-1]+ahalf[1:], -ahalf[1:-1]]
    A = (n+1)**2 * sparse.diags(diagonals, [-1, 0, 1], shape=(n,n), format='csr')
    return A

def relax(A, u, f, nu):
    n = A.shape[0]
    unew = u.copy()
    DE = sparse.tril(A, 0).tocsc()
    
    for i in range(nu):
        unew += sla.spsolve(DE, f - A * unew, permc_spec='NATURAL')

    return unew

def interpolation1d(nc, nf):
    """
    Linear interpolation
    """
    d = np.repeat([[1, 2, 1]], nc, axis=0).T
    I = np.zeros((3,nc))
    for i in range(nc):
        I[:,i] = [2*i, 2*i+1, 2*i+2]
    J = np.repeat([np.arange(nc)], 3, axis=0)
    P = sparse.coo_matrix(
        (d.ravel(), (I.ravel(), J.ravel()))
        ).tocsr()
    return 0.5 * P

def vcycle(Alist, Plist, u, f, nu):
    nlevels = len(Alist)
    ulist = [None for k in range(nlevels)]
    flist = [None for k in range(nlevels)]

    # down cycle
    for k in range(nlevels-1):
        Ak = Alist[k]
        Pk = Plist[k]
        u = relax(Ak, u, f, nu)
        ulist[k] = u
        flist[k] = f
        
        f = Pk.T * (f - Ak * u)
        u = np.zeros(f.shape)
        ulist[k+1] = u
        flist[k+1] = f

    # coarsest grid
    Ac = Alist[-1]
    flist[-1] = f
    ulist[-1] = sla.spsolve(Ac, f)

    # up cycle
    for k in range(nlevels-2, -1, -1):
        Ak = Alist[k]
        Pk = Plist[k]
        u = ulist[k]
        f = flist[k]
        uc = ulist[k+1]
        u += Pk * uc
        u = relax(Ak, u, f, nu)
    return u

kmax = 7
kmin = 2

# set up fine problem
n = 2**kmax - 1

xx = np.linspace(0, 1, n+2)[1:-1]

f = 2 - 12 * xx**2

ustar = xx**4 - xx**2
A = poissonop2(len(f))
udstar = sla.spsolve(A, f)

Alist = [A]
Plist = []
nf = n
while nf>3:
    nc = int((nf + 1) / 2 - 1)
    P = interpolation1d(nc, nf)
    Ac = P.T * Alist[-1] * P
    Plist.append(P)
    Alist.append(Ac)
    nf = nc

# set up smoothing sweeps
nu = 2

u = np.random.rand(len(f))
res = [hnorm(f - A * u)]
err = []

for i in range(12):
    u = vcycle(Alist, Plist, u, f, nu)
    res.append(hnorm(f - A * u))
    err.append(hnorm(u - ustar))

res = np.array(res)
resfactor = res[1:] / res[:-1]
print(resfactor)

[ 0.01172593  0.3827886   0.60717982  0.64626978  0.67354965  0.69168641
  0.70448037  0.71433598  0.72250966  0.7296079   0.73591318  0.74155601]

n

127