using LinearAlgebra, SparseArrays

# returns -∇² (discrete Laplacian, real-symmetric positive-definite) on n₁×n₂ grid
function ∇²(n₁,n₂)
    o₁ = ones(n₁)
    ∂₁ = spdiagm(n₁+1,n₁,-1=>-o₁,0=>o₁)
    o₂ = ones(n₂)
    ∂₂ = spdiagm(n₂+1,n₂,-1=>-o₂,0=>o₂)
    return kron(sparse(I,n₂,n₂), ∂₁'*∂₁) + kron(∂₂'*∂₂, sparse(I,n₁,n₁))
end

∇²(4,3)

dump(∇²(4,3))

Matrix(∇²(4,3))

∇²(10,5)

using PyPlot

spy(∇²(10,5))

using BenchmarkTools

A = ∇²(600,200) # a 120,000×120,000 matrix
b = rand(600*200)
@btime $A \ $b
@btime cholesky($A) \ $b
@btime lu($A) \ $b;

spy(sparse(cholesky(∇²(10,5)).L)')
title("sparse Cholesky factor, AMD ordering")

spy(sparse(cholesky(∇²(10,5), perm=1:50).L)')
title("sparse Cholesky factor, natural ordering")

N = 10:10:400
m = N.^2
loglog(m, [nnz(cholesky(∇²(n,n), perm=1:n*n)) for n in N], "r*")
loglog(m, [nnz(cholesky(∇²(n,n))) for n in N], "bo")
loglog(m, m.^(3/2), "r-")
loglog(m, m .* log2.(m), "b-")
xlabel(L"matrix size $m$")
ylabel("# nonzeros in Cholesky")
legend(["natural", "AMD", L"m^{3/2}", L"m \log_2(m)"], loc="upper left")

# return nested-dissection permutation for an m × n grid (in column-major order).  ndissect is an optional
# number of dissection steps before reverting to the normal ordering, and lm is an optional leading dimension
# (used for indexing computations).
function nest(m,n, ndissect=30, lm=m)
    if ndissect <= 0 || m*n < 5
        return vec([i + (j-1)*lm for i=1:m, j=1:n])
    elseif m >= n
        msep = div(m,2)
        N1 = nest(msep-1,n, ndissect-1, lm)
        N2 = nest(m-msep,n, ndissect-1, lm) .+ msep
        Ns = msep .+ (0:n-1)*lm
        return [N1; N2; Ns]
    else
        nsep = div(n,2)
        N1 = nest(m,nsep-1, ndissect-1, lm)
        N2 = nest(m,n-nsep, ndissect-1, lm) .+ nsep*lm
        Ns = (1:m) .+ (nsep-1)*lm
        return [N1; N2; Ns]
    end
end

sort(nest(30,60)) == [1:30*60;]

spy(sparse(cholesky(∇²(10,5), perm=nest(10,5)).L)')
title("sparse Cholesky factor, nested-dissection ordering")

m = 100
n = 120
nnz(cholesky(∇²(m,n), perm=1:m*n)), nnz(cholesky(∇²(m,n), perm=nest(m,n))),  nnz(cholesky(∇²(m,n)))

figure(figsize=(10,10))
n1, n2 = 10, 5
m = n1*n2

subplot(2,2,1)
F = cholesky(∇²(n1,n2), perm=1:m)
spy(sparse(F.L)')
title("natural, $(nnz(F)/m^2*100)% fill", y=1.1)

subplot(2,2,2)
F = cholesky(∇²(n1,n2), perm=nest(n1,n2, 1))
spy(sparse(F.L)')
title("1 dissection, $(nnz(F)/m^2*100)% fill", y=1.1)

subplot(2,2,3)
F = cholesky(∇²(n1,n2), perm=nest(n1,n2, 2))
spy(sparse(F.L)')
title("2 dissections, $(nnz(F)/m^2*100)% fill", y=1.1)

subplot(2,2,4)
F = cholesky(∇²(n1,n2), perm=nest(n1,n2, 3))
spy(sparse(F.L)')
title("3 dissections, $(nnz(F)/m^2*100)% fill", y=1.1)