Kernel polynomial method, take 2

Project.toml

@@ -6,6 +6,7 @@ version = "0.7.0"
 [deps]
 CrystalInfoFramework = "6007d9b0-c6b2-11e8-0510-1d10e825f3f1"
 DynamicPolynomials = "7c1d4256-1411-5781-91ec-d7bc3513ac07"
+ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 FilePathsBase = "48062228-2e41-5def-b9a4-89aafe57970f"
 HCubature = "19dc6840-f33b-545b-b366-655c7e3ffd49"
@@ -37,6 +38,7 @@ WGLMakiePrecompilesExt = "WGLMakie"
 [compat]
 CrystalInfoFramework = "0.6.0"
 DynamicPolynomials = "0.5.2"
+ElasticArrays = "1.2.12"
 FFTW = "1.4.5"
 FilePathsBase = "0.9.18"
 GLMakie = "0.9, 0.10"

examples/extra/KPM/chebyshev_usage.jl

@@ -0,0 +1,115 @@
+using Sunny, GLMakie, LinearAlgebra
+
+# Make the internal functions public for the purposes of illustration
+import Sunny: cheb_coefs, cheb_eval, apply_jackson_kernel!
+
+# Specify the function we wish to approximate and related parameters
+bose_reg(x, a) = 1/((2cosh(x/a))^a - 1)  # "Regularized" Bose function
+bounds = (-1.0, 1.0)  # Function domain
+nsamps = 4000  # Number of abscissa points used to sample the function (within bounds)
+npolys = 1000  # Number of polynomials to use in approximation
+
+# Calculate the coefficients
+a = 0.05
+coefs = cheb_coefs(npolys, nsamps, x -> bose_reg(x, a), bounds)
+
+# The coefficients may now be used to evaluate the approximant at arbitrary x
+# within the bounds
+x = 0.2
+y_exact = bose_reg(x, a)
+y_approx = cheb_eval(x, bounds, coefs)
+
+# Consider now an example of a function not defined on the interval [-1, 1].
+# We'll set the domain of our function to the interval [0, 10].
+bounds = (0.0, 10.0)
+
+# Define the function itself (a couple of Gaussian bumps)
+somefunc(x) = 0.5exp(-0.5((x - 2.0)/0.05)^2) + exp(-0.5((x - 7.0)/0.1)^2)
+
+# Calculate the coefficients
+coefs = cheb_coefs(npolys, nsamps, somefunc, bounds)
+
+# Construct the reference and approximation
+xs_ref = range(bounds[1], bounds[2], 2000)
+xs_rec = range(bounds[1], bounds[2], 400)
+ref = map(somefunc, xs_ref)
+rec = map(x -> cheb_eval(x, bounds, coefs), xs_rec)
+
+# Plot the results
+begin
+    fig = Figure()
+    ax = Axis(fig[1,1]; yscale=log10, ylabel=L"c_n", xlabel=L"n", title="Coefficients")
+    scatter!(ax, abs.(coefs); markersize=5.0)
+    ax = Axis(fig[1,2]; ylabel=L"f(x)", xlabel=L"x", title="Example Function")
+    lines!(ax, xs_ref, ref; color=(:black, 0.6), label="Reference")
+    scatter!(ax, xs_rec, rec; markersize=8.0, color=:red, label="Approximant")
+    axislegend(ax)
+    fig
+end
+
+
+
+# Examine approximations using different numbers of polynomials.
+begin
+    maxN = 40 # Number of polynomials to use in reconstruction -- choose between 2 and 1000
+
+    # Calculate the coefficients
+    a = 0.05
+    bounds = (-1, 1)
+    coefs = cheb_coefs(maxN, nsamps, x -> bose_reg(x, a), bounds)
+    coefs_jackson = apply_jackson_kernel!(copy(coefs))
+
+    # Create reference (evaluate original function on 1000 points)
+    xs_ref = range(bounds[1], bounds[2], 1000) 
+    ref = map(x -> bose_reg(x, 0.05), xs_ref) # "Ground truth"
+
+    # Reconstruct from coefficients, without and with Jackson kernel
+    xs_rec = range(bounds[1], bounds[2], 200)  # Points to evaluate reconstruction
+    rec = map(x -> cheb_eval(x, bounds, coefs), xs_rec)
+    jac = map(x -> cheb_eval(x, bounds, coefs_jackson), xs_rec)
+
+    # Plot results
+    p = lines(xs_ref, ref; color=(:black, 0.6), label="Reference")
+    scatter!(xs_rec, rec; marker=:circle, markersize=10.0, color=:orange, label="Reconstruction")
+    scatter!(xs_rec, jac; marker=:cross, markersize=10.0, color=:magenta, label="Reconstruction (Jackson)")
+    axislegend(p.axis)
+    p
+end
+
+
+# Example of how coefficients die off with varying width parameter
+begin
+    fig = Figure(resolution=(1200,500))
+    as = [0.1, 0.05, 0.025]  # Width parameter of regularized Bose function
+
+    # Plot coefficient magnitudes for each case
+    numtokeep = Int64[]
+    stored_coefs = []
+    for (n, a) in enumerate(as)
+        # Calculate and save coefficients
+        coefs = cheb_coefs(npolys, nsamps, x -> bose_reg(x, a), bounds)
+        push!(stored_coefs, coefs)
+
+        # Plot
+        axis = Axis(fig[1,n]; yscale=log10, ylabel = n == 1 ? L"c_n" : "", xlabel=L"n")
+        scatter!(axis, abs.(coefs); markersize=2.0)
+
+        # Figure out how many coefficients to use for reconstruction (clip when below certain value)
+        push!(numtokeep, findfirst(x -> abs(x) < 1e-5, coefs[begin:2:end]))
+    end
+
+    # Plot original function and reconstruction in each case
+    for (n, (maxM, coefs, a)) in enumerate(zip(numtokeep, stored_coefs, as))
+        # Reconstruct functions from saved coefficients -- used saved number of coefficients
+        rec = map(x -> cheb_eval(x, bounds, coefs[1:maxM]), xs_rec)
+
+        # Plot
+        ax = Axis(fig[2,n]; ylabel = n == 1 ? L"f(x)" : "", xlabel=L"x")
+        ref = map(x -> bose_reg(x, a), xs_ref) 
+        lines!(ax, xs_ref, ref; color=(:black, 0.6), label="Reference")
+        scatter!(ax, xs_rec, rec; marker=:circle, markersize=5.0, color=:orange, label="Reconstruction")
+        (n == 3) && Legend(fig[2,4], ax)
+    end
+
+    fig
+end

src/KPM/Chebyshev.jl

@@ -0,0 +1,58 @@
+@inline cheb_scale(x, bounds) = 2(x-bounds[1])/(bounds[2]-bounds[1]) - 1
+@inline cheb_unscale(x, bounds) = (x+1)*(bounds[2]-bounds[1])/2 + bounds[1]
+
+
+function apply_jackson_kernel!(coefs)
+    Mp = lastindex(coefs) + 2
+    for i in eachindex(coefs)
+        m = i-1
+        coefs[i] *= (1/Mp)*((Mp-m)cos(m*π/Mp) + sin(m*π/Mp)/tan(π/Mp))
+    end
+    return coefs
+end
+
+
+function cheb_coefs!(M, func, bounds; buf, plan)
+    N = length(buf)
+    for i in eachindex(buf)
+        x_i = cos((i-0.5)π / N)
+        buf[i] = func(cheb_unscale(x_i, bounds))
+    end
+
+    mul!(buf, plan, buf)
+    buf ./= N
+    buf[1] /= 2
+    return view(buf, 1:M)
+end
+
+"""
+    cheb_coefs(M, nsamples, func, bounds)
+
+Generate `M` coefficients of the Chebyshev expansion using `nsamples` of the
+function `func`. Sample points are taken within the interval specified by
+`bounds = (lo, hi)`.
+"""
+function cheb_coefs(M, nsamples, func, bounds)
+    @assert nsamples >= M
+    buf = zeros(nsamples)
+    plan = FFTW.plan_r2r!(buf, FFTW.REDFT10)
+    return cheb_coefs!(M, func, bounds; buf, plan)
+end
+
+"""
+    cheb_eval(x, bounds, coefs)
+
+Evaluate a function, specified in terms of the Chebyshev `coefs`, at point `x`.
+`bounds` specifies the domain of the function.
+"""
+function cheb_eval(x, bounds, coefs)
+    x = cheb_scale(x, bounds)
+    Tn2, Tn1 = 1, x
+    ret = coefs[1]*Tn2 + coefs[2]*Tn1
+    for coef in @view coefs[3:end]
+        Tn3 = 2x*Tn1 - Tn2
+        ret += coef*Tn3
+        Tn1, Tn2 = Tn3, Tn1
+    end
+    return ret
+end

src/KPM/Lanczos.jl

@@ -0,0 +1,79 @@
+
+"""
+    parallel_lanczos!(buf1, buf2, buf3; niters, mat_vec_mul!, inner_product!)
+
+Parallelized Lanczos. Takes functions for matrix-vector multiplication and inner
+products. Returns a list of tridiagonal matrices. `niters` sets the number of
+iterations used by the Lanczos algorithm.
+
+References:
+- [H. A. van der Vorst, Math. Comp. 39, 559-561
+  (1982)](https://doi.org/10.1090/s0025-5718-1982-0669648-0).
+- [M. Grüning, A. Marini, X. Gonze, Comput. Mater. Sci. 50, 2148-2156
+  (2011)](https://doi.org/10.1016/j.commatsci.2011.02.021).
+"""
+function parallel_lanczos!(vprev, v, w;  niters, mat_vec_mul!, inner_product)
+    # ...
+end
+
+
+function lanczos(swt, q_reshaped, niters)
+    L = nbands(swt)
+
+    function inner_product(w, v)
+        return @views real(dot(w[1:L], v[1:L]) - dot(w[L+1:2L], v[L+1:2L]))
+    end
+
+    function mat_vec_mul!(w, v, q_reshaped)
+        mul_dynamical_matrix!(swt, reshape(w, 1, :), reshape(v, 1, :), [q_reshaped])
+        view(w, L+1:2L) .*= -1
+    end
+
+    αs = zeros(Float64, niters)   # Diagonal part
+    βs = zeros(Float64, niters-1) # Off-diagonal part
+
+    vprev = randn(ComplexF64, 2L)
+    v = zeros(ComplexF64, 2L)
+    w = zeros(ComplexF64, 2L)
+
+    mat_vec_mul!(w, vprev, q_reshaped)
+    b0 = sqrt(inner_product(vprev, w))
+    vprev = vprev / b0
+    w = w / b0
+    αs[1] = inner_product(w, w)
+    @. w = w - αs[1]*vprev
+    for j in 2:niters
+        @. v = w
+        mat_vec_mul!(w, v, q_reshaped)
+        βs[j-1] = sqrt(inner_product(v, w)) # maybe v?
+        if abs(βs[j-1]) < 1e-12
+            # TODO: Terminate early if all β are small
+            βs[j-1] = 0
+            @. v = w = 0
+        else
+            @. v = v / βs[j-1]
+            @. w = w / βs[j-1]
+        end
+        αs[j] = inner_product(w, w)
+        @. w = w - αs[j]*v - βs[j-1]*vprev
+        v, vprev = vprev, v
+    end
+
+    return SymTridiagonal(αs, βs)
+end
+
+
+"""
+    eigbounds(A, niters; extend=0.0)
+
+Returns estimates of the extremal eigenvalues of Hermitian matrix `A` using the
+Lanczos algorithm. `niters` should be given a value smaller than the dimension
+of `A`. `extend` specifies how much to shift the upper and lower bounds as a
+percentage of the total scale of the estimated eigenvalues.
+"""
+function eigbounds(swt, q_reshaped, niters; extend)
+    tridiag = lanczos(swt, Vec3(q_reshaped), niters)
+    lo, hi = extrema(eigvals!(tridiag)) # TODO: pass irange=1 and irange=2L
+    slack = extend*(hi-lo)
+    return lo-slack, hi+slack
+end