KPM Implementation.

Project.toml

@@ -23,6 +23,7 @@ ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 RowEchelon = "af85af4c-bcd5-5d23-b03a-a909639aa875"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 Spglib = "f761d5c5-86db-4880-b97f-9680a7cccfb5"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

examples/extra/chebyshev_usage.jl

@@ -0,0 +1,124 @@
+using Sunny, GLMakie, LinearAlgebra
+
+# Make the internal functions public for the purposes of illustration
+import Sunny: cheb_coefs, cheb_eval, 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)
+
+# If we wish to apply the Jackson kernel to the coefficients, this may be passed
+# as an additional argument.
+kernel = jackson_kernel(npolys)
+y_approx = cheb_eval(x, bounds, coefs; kernel)
+
+# 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 = 20 # Number of polynomials to use in reconstruction -- choose between 2 and 1000
+
+    # Calculate the coefficients
+    a = 0.05
+    bounds = (-1, 1)
+    coefs = cheb_coefs(npolys, nsamps, x -> bose_reg(x, a), bounds)
+
+    # 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
+    kernel = jackson_kernel(maxN)  # Coefficient weightings from Jackson kernel
+
+    rec = map(x -> cheb_eval(x, bounds, coefs; maxN), xs_rec)
+    jac = map(x -> cheb_eval(x, bounds, coefs; kernel, maxN), 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[0:2:end]))
+    end
+
+    # Plot original function and reconstruction in each case
+    for (n, (maxN, coefs, a)) in enumerate(zip(numtokeep, stored_coefs, as))
+        # Reconstruct functions from saved coefficients -- used saved number of coefficients
+        kernel = jackson_kernel(maxN)
+        rec = map( x -> cheb_eval(x, bounds, coefs; maxN), xs_rec)
+        jac = map( x -> cheb_eval(x, bounds, coefs; kernel, maxN), 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")
+        scatter!(ax, xs_rec, jac; marker=:cross, markersize=5.0, color=:magenta, label="Reconstruction (Jackson)")
+        (n == 3) && Legend(fig[2,4], ax)
+    end
+
+    fig
+end

src/SpinWaveTheory/Chebyshev.jl

@@ -0,0 +1,68 @@
+"""
+    jackson_kernel(N)
+
+Generate a Jackson kernel for `N` Chebyshev coefficients. For use with
+`cheb_eval`.
+"""
+function jackson_kernel(N)
+    Np = N + 1.0
+    map(n -> (1/Np)*((Np-n)cos(n*π/Np) + sin(n*π/Np)/tan(π/Np)), 0:N-1) |>
+        A -> OffsetArray(A, 0:N-1)
+end
+@inline scale(x, bounds) = 2(x-bounds[1])/(bounds[2]-bounds[1]) - 1
+@inline unscale(x, bounds) = (x+1)*(bounds[2]-bounds[1])/2 + bounds[1]
+
+
+# Add FFT planned version of following. Can make fully non-allocating version
+# when it's time to optimize (e.g., remove internal buffer allocation)
+"""
+    cheb_coefs(N, Nsamp, func, bounds)
+
+Generate `N` coefficients of the Chebyshev expansion using `Nsamp` samples of
+the function `func`. Sample points are taken within the interval specified by
+`bounds`, e.g., `bounds=(-1,1)`.
+"""
+function cheb_coefs(N, Nsamp, func, bounds)
+    @assert Nsamp >= N
+    buf = OffsetArray(zeros(Nsamp), 0:Nsamp-1) 
+    out = OffsetArray(zeros(N),  0:N-1)
+    for i in 0:Nsamp-1
+        x_i = cos((i+0.5)π / Nsamp)
+        buf[i] = func(unscale(x_i, bounds))
+    end
+    FFTW.r2r!(buf.parent, FFTW.REDFT10)
+    for n in 0:N-1
+        out[n] = (iszero(n) ? 0.5 : 1.0) * buf[n] / Nsamp
+    end
+    return out
+end
+
+
+"""
+    cheb_eval(x, bounds, coefs::OffsetArray; kernel=nothing, maxN = nothing)
+
+Evaluate a function, specified in terms of the Chebyshev `coefs`, at point `x`.
+`bounds` specifies the domain of the function. `kernel` is used to specify a
+weighting of the `coefs`, using, for example, `jackson_kernel`. `maxN` specifies
+how many coefficients (polynomials) to retain in the reconstruction.
+"""
+function cheb_eval(x, bounds, coefs::T; kernel = nothing, maxN = nothing) where T <: OffsetArray
+    ncoefs = length(coefs)
+    @assert ncoefs > 2
+    maxN = isnothing(maxN) ? ncoefs : maxN
+    @assert maxN <= ncoefs
+    g = isnothing(kernel) ? _ -> 1.0 : n -> kernel[n]
+
+    return cheb_eval_aux(scale(x, bounds), coefs, g, maxN)
+end
+
+function cheb_eval_aux(x, c, g, maxN)
+    Tn2, Tn1 = 1, x
+    out = g(0)*c[0]*Tn2 + g(1)*c[1]*Tn1
+    for n in 2:maxN-1
+        Tn = 2x*Tn1 - Tn2
+        out += g(n)*c[n]*Tn 
+        Tn1, Tn2 = Tn, Tn1
+    end
+    out
+end

src/SpinWaveTheory/Lanczos.jl

@@ -37,21 +37,20 @@ end
 function lanczos_aux!(αs, βs, buf, A, niters)
     v, vprev, w = view(buf,:,1), view(buf,:,2), view(buf,:,3)
 
-    @label initialize_lanczos
-    randn!(v)
-    normalize!(v)
-    mul!(w, A, v)
-    αs[1] = real(w⋅v)
-    @. w = w - αs[1]*v
+    randn!(vprev)
+    normalize!(vprev)
+    mul!(w, A, vprev)
+    αs[1] = real(w⋅vprev)
+    @. w = w - αs[1]*vprev
 
     for j in 2:niters
-        v, vprev = vprev, v
         βs[j-1] = norm(w)
-        (abs(βs[j-1]) < 1e-12) && (@goto initialize_lanczos)
+        iszero(βs[j-1]) && return
         @. v = w/βs[j-1]
         mul!(w, A, v)
         αs[j] = real(w⋅v)
         @. w = w - αs[j]*v - βs[j-1]*vprev
+        v, vprev = vprev, v
     end
 
     return nothing