diff --git a/Project.toml b/Project.toml index 35cfdb58c..1ea885f5a 100644 --- a/Project.toml +++ b/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" diff --git a/examples/extra/chebyshev_usage.jl b/examples/extra/chebyshev_usage.jl new file mode 100644 index 000000000..ff4e42bf1 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/SpinWaveTheory/Chebyshev.jl b/src/SpinWaveTheory/Chebyshev.jl new file mode 100644 index 000000000..3cadb6514 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/SpinWaveTheory/Lanczos.jl b/src/SpinWaveTheory/Lanczos.jl index 8d9fa083d..ef45cb3e6 100644 --- a/src/SpinWaveTheory/Lanczos.jl +++ b/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 diff --git a/src/SpinWaveTheory/kpm.jl b/src/SpinWaveTheory/kpm.jl index b9297ea0f..4aaa389ce 100644 --- a/src/SpinWaveTheory/kpm.jl +++ b/src/SpinWaveTheory/kpm.jl @@ -1,45 +1,163 @@ -struct KPMIntensityFormula{T} - P :: Int64 - kT :: Float64 - σ :: Float64 - broadening - kernel - string_formula :: String - calc_intensity :: Function -end +""" +This code will implement the Kernal Polynomial method to approximate the dynamical spin structure factor (DSSF). +The method takes advantage of the fact that the DSSF can be written as a matrix function and hence we can avoid +diagonalizing the Hamiltonian. Futher, implementing the KPM we do not even need to explicity construct the matrix +function, instead we only need to evaluate the application of the function to some vector. For sparce matrices this +is efficient making this approximate approach useful in systems with large unit cells. +""" + +""" + bose_function(kT, x) -function Base.show(io::IO, formula::KPMIntensityFormula{T}) where T - print(io,"KPMIntensityFormula{$T}") +Returns the Bose occupation factor for energy, x, at temperature kT. +""" + +function bose_function(kT, x) + return 1 / (exp(x/kT) - 1) end -function Base.show(io::IO, ::MIME"text/plain", formula::KPMIntensityFormula{T}) where T - printstyled(io, "Quantum Scattering Intensity Formula (KPM Method)\n";bold=true, color=:underline) - formula_lines = split(formula.string_formula,'\n') +""" + get_all_coefficients(M, ωs, broadening, σ, kT,γ) + +Retrieves the Chebyshev coefficients up to index, M for a user-defined lineshape. A numerical regularization is applied to +treat the divergence of the dynamical correlation function at small energy. Here, σ² represents an energy cutoff scale +which should be related to the energy resolution. γ is the maximum eigenvalue used to rescale the spectrum to lie on the +interval [-1,1]. - intensity_equals = " Intensity(Q,ω) = " - println(io,"At any (Q,ω), with S = ...:") - println(io) - println(io,intensity_equals,formula_lines[1]) - for i = 2:length(formula_lines) - precursor = repeat(' ', textwidth(intensity_equals)) - println(io,precursor,formula_lines[i]) +""" +function get_all_coefficients(M, ωs, broadening, σ, kT,γ) + f(ω, x) = tanh((x/σ)^2) * broadening(ω, x*γ, σ) * (1 + bose_function(kT, x)) + output = OffsetArray(zeros(M, length(ωs)), 0:M-1, 1:length(ωs)) + for i in eachindex(ωs) + output[:, i] = cheb_coefs(M, 2M, x -> f(ωs[i], x), (-1, 1)) end - println(io,"P = $(formula.P), kT = $(formula.kT), σ = $(formula.σ)") + return output end +""" + apply_kernel(offset_array, kernel, M) -function intensity_formula_kpm(f::Function,swt::SpinWaveTheory,corr_ix::AbstractVector{Int64}; P =, return_type = Float64, string_formula = "f(Q,ω,S{α,β}[ix_q,ix_ω])") +Applies the Jackson kernel (defined in Chebyshev.jl) to an OffsetArray. The Jackson kernel damps the coefficients of the +Chebyshev expansion to reduce "Gibbs oscillations" or "ringing" -- numerical artifacts present due to the truncation offset_array +the Chebyshev series. It should be noted that employing the Jackson kernel comes at a cost to the energy resolution. +""" - error("KPM not yet implemented") +# +function apply_kernel(offset_array, kernel, M) + kernel === "jackson" ? offset_array .= offset_array .* jackson_kernel(M) : + return offset_array +end + +""" + kpm_dssf(swt::SpinWaveTheory, qs,ωlist,P::Int64,kT,σ,broadening; kernel) - stuff = setup_stuff(swt) - formula = function(swt::SpinWaveTheory,q::Vec3,ω::Float64) - Sαβ = do_KPM(swt,stuff) - k = swt.recipvecs_chem * q - intensity = f(k,ω,Sαβ[corr_ix]) +Calculated the Dynamical Spin Structure Factor (DSSF) using the Kernel Polynomial Method (KPM). Requires input in the form of a +SpinWaveTheory which contains System information and the rotated spin matrices. The calculation is carried out at each wavevectors +in qs for all energies appearing in ωlist. The Chebyshev expansion is taken to P terms and the lineshape is specified by the user- +defined function, broadening. kT is required for the calcualation of the bose function and σ is the width of the lineshape and +defines the low energy cutoff σ². There is a keyword argument, kernel, which speficies a damping kernel. +""" + +function kpm_dssf(swt::SpinWaveTheory, qs,ωlist,P::Int64,kT,σ,broadening; kernel) + # P is the max Chebyshyev coefficient + (; sys, s̃_mat, T̃_mat, Q̃_mat,positions_chem) = swt + qs = Sunny.Vec3.(qs) + Nm, Ns = length(sys.dipoles), sys.Ns[1] # number of magnetic atoms and dimension of Hilbert space + Nf = sys.mode == :SUN ? Ns-1 : 1 + N=Nf+1 + nmodes = Nf*Nm + M = sys.mode == :SUN ? 1 : (Ns-1) # scaling factor (=1) if in the fundamental representation + sqrt_M = √M #define prefactors + sqrt_Nm_inv = 1.0 / √Nm #define prefactors + Ĩ = spdiagm([ones(nmodes); -ones(nmodes)]) + n_iters = 50 + Hmat = zeros(ComplexF64, 2*nmodes, 2*nmodes) + Avec_pref = zeros(ComplexF64, Nm) # initialize array of some prefactors + chebyshev_moments = OffsetArray(zeros(ComplexF64,3,3,length(qs),P),1:3,1:3,1:length(qs),0:P-1) + Sαβs = zeros(ComplexF64,3,3,length(qs),length(ωlist)) + for qidx in CartesianIndices(qs) + q = qs[qidx] + _, qmag = Sunny.chemical_to_magnetic(swt, q) + u = zeros(ComplexF64,3,2*nmodes) + if sys.mode == :SUN + swt_hamiltonian_SUN!(swt, qmag, Hmat) + else + #swt_hamiltonian_dipole!(swt, qmag, Hmat) # this will break, update with new dipole mode code later + throw("Please set mode = :SUN ") + end + D = 2.0*sparse(Hmat) # calculate D (factor of 2 for correspondence) + lo,hi = Sunny.eigbounds(D,n_iters; extend=0.25) # calculate bounds + γ=max(lo,hi) # select upper bound (combine with the preceeding line later) + A = Ĩ*D / γ + # u(q) calculation) + for site = 1:Nm + # note that d is the chemical coordinates + chemical_coor = positions_chem[site] # find chemical coords + phase = exp(2*im * π * dot(q, chemical_coor)) # calculate phase + Avec_pref[site] = sqrt_Nm_inv * phase * sqrt_M # define the prefactor of the tS matrices + end + # calculate u(q) + for site=1:Nm + @views tS_μ = s̃_mat[:, :, :, site]*Avec_pref[site] + for μ=1:3 + for j=2:N + u[μ,(j-1)+(site-1)*(N-1) ]=tS_μ[j,1,μ] + u[μ,(N-1)*Nm+(j-1)+(site-1)*(N-1) ]=tS_μ[1,j,μ] + end + end + end + for β=1:3 + α0 = zeros(ComplexF64,2*nmodes) + α1 = zeros(ComplexF64,2*nmodes) + mul!(α0,Ĩ,u[β,:]) # calculate α0 + mul!(α1,A,α0) # calculate α1 + for α=1:3 + chebyshev_moments[α,β,qidx,0] = real(dot(u[α,:],α0)) + chebyshev_moments[α,β,qidx,1] = real(dot(u[α,:],α1)) + end + for m=2:P-1 + αnew = zeros(ComplexF64,2*nmodes) + mul!(αnew,A,α1) + @. αnew = 2*αnew - α0 + for α=1:3 + chebyshev_moments[α,β,qidx,m] = real(dot(u[α,:],αnew)) + end + (α1, α0) = (αnew, α1) + end + end + ωdep = get_all_coefficients(P,ωlist,broadening,σ,kT,γ) + apply_kernel(ωdep,kernel,P) + for w=1:length(ωlist) + for α=1:3 + for β=1:3 + Sαβs[α,β,qidx,w] = sum(chebyshev_moments[α,β,qidx,:] .* ωdep[:,w]) + end + end + end end - KPMIntensityFormula{return_type}(P,kT,σ,broadening,kernel,string_formula,formula) + return Sαβs end +""" + kpm_intensities(swt::SpinWaveTheory, qs,ωlist,P::Int64,kT,σ,broadening; kernel) +Calculated the neutron scattering intensity using the Kernel Polynomial Method (KPM). Calls KPMddsf and so takes the same parameters. +Requires input in the form of a SpinWaveTheory which contains System information and the rotated spin matrices. The calculation is +carried out at each wavevectors in qs for all energies appearing in ωlist. The Chebyshev expansion is taken to P terms and the +lineshape is specified by the user-defined function, broadening. kT is required for the calcualation of the bose function and σ is +the width of the lineshape and defines the low energy cutoff σ². There is an optional keyword argument, kernel, which speficies a +damping kernel. The default is to include no damping. +""" +function kpm_intensities(swt::SpinWaveTheory, qs, ωvals,P::Int64,kT,σ,broadening; kernel = nothing) + (; sys) = swt + qs = Sunny.Vec3.(qs) + Sαβs = kpm_dssf(swt, qs,ωvals,P,kT,σ,broadening;kernel) + num_ω = length(ωvals) + is = zeros(Float64, size(qs)..., num_ω) + for qidx in CartesianIndices(qs) + polar_mat = Sunny.polarization_matrix(swt.recipvecs_chem * qs[qidx]) + is[qidx, :] = real(sum(polar_mat .* Sαβs[:,:,qidx,:],dims=(1,2))) + end + return is +end diff --git a/src/Sunny.jl b/src/Sunny.jl index 846749472..80851dfa2 100644 --- a/src/Sunny.jl +++ b/src/Sunny.jl @@ -27,6 +27,9 @@ import Colors: distinguishable_colors, RGB, Colors import Inflate: inflate_gzip import Random: randstring, RandomDevice +# Specific to KPM +import SparseArrays: spzeros, sparse, spdiagm + const Vec3 = SVector{3, Float64} const Mat3 = SMatrix{3, 3, Float64, 9} const CVec{N} = SVector{N, ComplexF64} @@ -82,7 +85,9 @@ export minimize_energy! include("SpinWaveTheory/SpinWaveTheory.jl") include("SpinWaveTheory/SWTCalculations.jl") include("SpinWaveTheory/Lanczos.jl") -export SpinWaveTheory, dispersion, dssf, delta_function_kernel +include("SpinWaveTheory/Chebyshev.jl") +include("SpinWaveTheory/KPM.jl") +export SpinWaveTheory, dispersion, intensities, dssf, kpm_dssf, kpm_intensities include("SampledCorrelations/SampledCorrelations.jl") include("SampledCorrelations/CorrelationUtils.jl") diff --git a/test/test_lswt.jl b/test/test_lswt.jl index ed9d1d613..6b35a033c 100644 --- a/test/test_lswt.jl +++ b/test/test_lswt.jl @@ -10,6 +10,55 @@ @test (abs(lo/vals[1] - 1) < 0.025) && (abs(hi/vals[end] - 1) < 0.025) end +@testitem "Chebyshev Approximations" begin + using LinearAlgebra + + function cheb_coefs_manual(M, Mq, func, bounds) + function cheb_array!(out, x) + len = length(out) + @assert len > 2 + out[0], out[1] = 1.0, x + for m = 2:len-1 + out[m] = 2x*out[m-1] - out[m-2] + end + return nothing + end + + @assert Mq > M + out = Sunny.OffsetArray(zeros(M), 0:M-1) # To avoid bringing OffsetArrays into testing environment + fp = Sunny.OffsetArray(zeros(M), 0:M-1) + T = Sunny.OffsetArray(zeros(M), 0:M-1) + for i in 0:Mq-1 + x_i = cos((i+0.5)π/Mq) + f_i = func(Sunny.unscale(x_i, bounds)) + cheb_array!(T, x_i) + @. fp += f_i * T + end + for m in 0:M-1 + out[m] = (iszero(m) ? 1.0 : 2.0 ) * fp[m] / Mq + end + return out + end + + # Set up test function + bounds = (-1, 1) + σ = 0.1 + func(x) = (1/√(2π*σ^2))*exp(-0.5*(x/σ)^2) + + # Check coefficients from Sunny function match manual calculation + M = 100 + Mq = 400 + coefs = Sunny.cheb_coefs(M, Mq, func, bounds) + coefs_ref = cheb_coefs_manual(M, Mq, func, bounds) + @test norm(coefs - coefs_ref) < 1e-12 + + # Check resulting approximation sufficiently close to ground truth + xs = -0.9:0.1:0.9 + ref = func.(xs) + rec = map(x -> Sunny.cheb_eval(x, bounds, coefs), xs) + @test norm(rec - ref) < 1e-12 +end + @testitem "Single Ion" begin J, J′, D = 1.0, 0.1, 5.0