Add KPM to main

Project.toml

@@ -17,6 +17,7 @@ Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 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"
@@ -44,6 +45,7 @@ Optim = "1.7.6"
 Printf = "1.9"
 Random = "1.9"
 RowEchelon = "0.2.1"
+SparseArrays = "1.9"
 SpecialFunctions = "2.1.0"
 Spglib = "0.8.0"
 StaticArrays = "1.3.4"

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/Intensities/ElementContraction.jl

@@ -187,3 +187,15 @@ function intensity_formula(sc::SampledCorrelations, contractor::Contraction{T};
         intensity = contract(correlations, k, contractor)
     end
 end
+
+function intensity_formula_kpm(swt::SpinWaveTheory, mode::Symbol; kwargs...)
+    contractor, string_formula = contractor_from_mode(swt, mode)
+    intensity_formula_kpm(swt, contractor; string_formula, kwargs...)
+end
+
+function intensity_formula_kpm(swt::SpinWaveTheory, contractor::Contraction{T}; kwargs...) where T
+    intensity_formula_kpm(swt,required_correlations(contractor); return_type = T,kwargs...) do k,correlations
+        intensity = contract(correlations, k, contractor)
+    end
+end
+

src/SpinWaveTheory/Chebyshev.jl

@@ -0,0 +1,83 @@
+"""
+    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
+
+function view_cheb_samps(N, Nsamp, func, bounds)
+    @assert Nsamp >= N
+    xx = OffsetArray(zeros(Nsamp), 0:Nsamp-1) 
+    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)
+        xx[i] = unscale(x_i, bounds)
+        buf[i] = func(unscale(x_i, bounds))
+    end
+    xx, buf
+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/DispersionAndIntensities.jl

@@ -375,4 +375,4 @@ function intensity_formula(f::Function,swt::SpinWaveTheory,corr_ix::AbstractVect
     end
     output_type = isnothing(kernel) ? BandStructure{nmodes,return_type} : return_type
     SpinWaveIntensityFormula{output_type}(string_formula, kernel_edep, calc_intensity)
-end
+end

src/SpinWaveTheory/HamiltonianDipole.jl

@@ -242,4 +242,108 @@ function multiply_by_hamiltonian_dipole_aux!(y, x, phasebuf, qphase, swt)
     @inbounds @. y += swt.energy_ϵ * x
 
     nothing
-end
+end
+
+# This is a stopgap measure to avoid unnecessary reshaping. Revisit on merge.
+
+# Calculate y = H*x, where H is the quadratic Hamiltonian matrix (dynamical matrix).
+function multiply_by_hamiltonian_dipole!(y, x, swt::SpinWaveTheory, q_reshaped::Vec3)
+    (; sys, data) = swt
+    (; local_rotations, stevens_coefs) = data
+
+    N = swt.sys.Ns[1]
+    S = (N-1)/2
+    L = nbands(swt)
+
+    x = reshape(x, natoms(sys.crystal), 2)
+    y = reshape(y, natoms(sys.crystal), 2)
+    y .= 0
+
+    # Add Zeeman term
+    (; extfield, gs, units) = sys
+    for i in 1:L
+        B = units.μB * (gs[1, 1, 1, i]' * extfield[1, 1, 1, i]) 
+        B′ = dot(B, local_rotations[i][:, 3]) / 2 
+
+        y[i, 1] += B′ * x[i, 1]
+        y[i, 2] += B′ * x[i, 2]
+    end
+
+    # Add pairwise terms 
+    for ints in sys.interactions_union
+
+        # Bilinear exchange
+        for coupling in ints.pair
+            (; isculled, bond) = coupling
+            isculled && break
+            i, j = bond.i, bond.j
+            phase = exp(2π*im * dot(q_reshaped, bond.n)) # Phase associated with periodic wrapping
+
+            if !iszero(coupling.bilin)
+                J = coupling.bilin  # This is Rij in previous notation (transformed exchange matrix)
+
+                P = 0.25 * (J[1, 1] - J[2, 2] - im*J[1, 2] - im*J[2, 1])
+                Q = 0.25 * (J[1, 1] + J[2, 2] - im*J[1, 2] + im*J[2, 1])
+
+                y[i, 1] += Q * phase * x[j, 1]
+                y[j, 1] += conj(Q) * conj(phase) * x[i, 1]
+                y[i, 2] += conj(Q) * phase * x[j, 2]
+                y[j, 2] += Q * conj(phase) * x[i, 2]
+
+                y[i, 2] += P * phase * x[j, 1]
+                y[j, 2] += P * conj(phase) * x[i, 1]
+                y[i, 1] += conj(P) * phase * x[j, 2]
+                y[j, 1] += conj(P) * conj(phase) * x[i, 2]
+
+                y[i, 1] -= 0.5 * J[3, 3] * x[i, 1]
+                y[j, 1] -= 0.5 * J[3, 3] * x[j, 1]
+                y[i, 2] -= 0.5 * J[3, 3] * x[i, 2]
+                y[j, 2] -= 0.5 * J[3, 3] * x[j, 2]
+            end
+
+            # Biquadratic exchange
+            if !iszero(coupling.biquad)
+                J = coupling.biquad  # Transformed quadrupole exchange matrix
+
+                y[i, 1] += -6J[3, 3] * x[i, 1]
+                y[j, 1] += -6J[3, 3] * x[j, 1]
+
+                y[i, 2] += -6J[3, 3] * x[i, 2]
+                y[j, 2] += -6J[3, 3] * x[j, 2]
+
+                y[i, 2] += 12*(J[1, 3] - im*J[5, 3]) * x[i, 1]
+                y[i, 1] += 12*(J[1, 3] + im*J[5, 3]) * x[i, 2]
+                y[j, 2] += 12*(J[3, 1] - im*J[3, 5]) * x[j, 1]
+                y[j, 1] += 12*(J[3, 1] + im*J[3, 5]) * x[j, 2]
+
+                P = 0.25 * (-J[4, 4]+J[2, 2] - im*( J[4, 2]+J[2, 4]))
+                Q = 0.25 * ( J[4, 4]+J[2, 2] - im*(-J[4, 2]+J[2, 4]))
+
+                y[i, 1] += Q * phase * x[j, 1]
+                y[j, 1] += conj(Q) * conj(phase) * x[i, 1]
+                y[i, 2] += conj(Q) * phase * x[j, 2]
+                y[j, 2] += Q * conj(phase) * x[i, 2]
+
+                y[i, 2] += P * phase * x[j, 1]
+                y[j, 2] += P * conj(phase) * x[i, 1]
+                y[i, 1] += conj(P) * phase * x[j, 2]
+                y[j, 1] += conj(P) * conj(phase) * x[i, 2]
+            end
+        end
+    end
+
+    # Add single-ion anisotropy
+    for i in 1:L
+        (; c2, c4, c6) = stevens_coefs[i]
+        y[i, 1] += (-3S*c2[3] - 40*S^3*c4[5] - 168*S^5*c6[7]) * x[i, 1]
+        y[i, 2] += (-3S*c2[3] - 40*S^3*c4[5] - 168*S^5*c6[7]) * x[i, 2]
+        y[i, 1] += (-im*(S*c2[5] + 6S^3*c4[7] + 16S^5*c6[9]) + (S*c2[1] + 6S^3*c4[3] + 16S^5*c6[5])) * x[i, 2]
+        y[i, 2] += (im*(S*c2[5] + 6S^3*c4[7] + 16S^5*c6[9]) + (S*c2[1] + 6S^3*c4[3] + 16S^5*c6[5])) * x[i, 1]
+    end
+
+    # Add small constant shift for positive-definiteness
+    @. y += swt.energy_ϵ * x
+
+    nothing
+end
+