Fix sampled correlations tests

src/Intensities/Interpolation.jl

@@ -2,9 +2,6 @@ abstract type InterpolationScheme{NumInterp} end
 struct NoInterp <: InterpolationScheme{1} end
 struct LinearInterp <: InterpolationScheme{8} end
 
-#=
-ddtodo: Explanation of interpolation "API"
-=# 
 
 function interpolated_intensity(::SampledCorrelations, _, _, stencil_intensities, ::NoInterp) 
     return only(stencil_intensities)
@@ -35,17 +32,10 @@ end
 
 
 function stencil_points(sc::SampledCorrelations, q, ::NoInterp)
-
-    # Each of the following lines causes a 32 byte allocation
     Ls = size(sc.samplebuf)[2:4] 
     m = round.(Int, Ls .* q)
     im = map(i -> mod(m[i], Ls[i])+1, (1, 2, 3)) |> CartesianIndex{3}
 
-    ## The following lines cause no allocations, but don't seem to be any faster.
-    #     _, L1, L2, L3, _, _ = size(sc.samplebuf)
-    #     m = (round(Int, L1*q[1]), round(Int, L2*q[2]), round(Int, L3*q[3]))
-    #     im = CartesianIndex{3}(mod(m[1], L1)+1, mod(m[2], L2)+1, mod(m[3], L3)+1)
-
     return (m,), (im,)
 end
 
@@ -140,21 +130,13 @@ function intensities_interpolated(sc::SampledCorrelations, qs;
     formula = intensity_formula(sc,:perp) :: ClassicalIntensityFormula,
     interpolation = :round,
     negative_energies = false,
-    static_warn = true
+    instantaneous_warning = true
 )
     qs = Vec3.(qs)
-
-    # If working on reshaped system, assume qs given as coordinates in terms of
-    # reciprocal vectors of original crystal and convert them to qs in terms of
-    # the reciprocal vectors of the reshaped crystal.
-    if !isnothing(sc.origin_crystal)
-        rvecs_reshaped = inv(sc.crystal.latvecs)'       # Note, leading 2π will cancel
-        rvecs_origin = inv(sc.origin_crystal.latvecs)'
-        qs = map(q -> rvecs_reshaped \ rvecs_origin * q, qs)
-    end
+    qs = map!(q -> sc.crystal.recipvecs \ q, qs, qs) 
 
     # Make sure it's a dynamical structure factor 
-    if static_warn && size(sc.data, 7) == 1
+    if instantaneous_warning && size(sc.data, 7) == 1
         error("`intensities_interpolated` given a SampledCorrelations with no dynamical information. Call `instant_intensities_interpolated` to retrieve instantaneous (static) structure factor data.")
     end
 
@@ -174,7 +156,7 @@ function intensities_interpolated(sc::SampledCorrelations, qs;
     return_type = typeof(formula).parameters[1]
     intensities = zeros(return_type, size(qs)..., nω)
 
-    # Call type stable version of the function
+    # Call type-stable version of the function
     intensities_interpolated!(intensities, sc, qs, ωvals, interp, formula, stencil_info, return_type)
 
     return intensities
@@ -211,7 +193,7 @@ i.e., ``𝒮(𝐪,ω)``, the ``ω`` information is integrated out.
 function instant_intensities_interpolated(sc::SampledCorrelations, qs; kwargs...)
     datadims = size(qs)
     ndims = length(datadims)
-    vals = intensities_interpolated(sc, qs; static_warn=false, kwargs...)
+    vals = intensities_interpolated(sc, qs; instantaneous_warning=false, kwargs...)
     static_vals = sum(vals, dims=(ndims+1,))
     return reshape(static_vals, datadims)
 end

src/SampledCorrelations/CorrelationUtils.jl

@@ -12,7 +12,7 @@ function all_exact_wave_vectors(sc::SampledCorrelations; bzsize=(1,1,1))
     lo = map(L -> 1 - div(L, 2), up) .- offsets
     qs = zeros(Vec3, up...)
     for (k, lz) in enumerate(lo[3]:hi[3]), (j, ly) in enumerate(lo[2]:hi[2]), (i, lx) in enumerate(lo[1]:hi[1])
-        qs[i,j,k] = Vec3(lx/Ls[1], ly/Ls[2], lz/Ls[3]) 
+        qs[i,j,k] = sc.crystal.recipvecs * Vec3(lx/Ls[1], ly/Ls[2], lz/Ls[3]) 
     end
     return qs
 end
@@ -36,7 +36,4 @@ function ωs(sc::SampledCorrelations; negative_energies=false)
         ωvals[i] -= 2ωvals[hω]
     end
     return negative_energies ? ωvals : ωvals[1:hω]
-end
-
-
-orig_crystal(sc::SampledCorrelations) = isnothing(sc.origin_crystal) ? sc.crystal : sc.origin_crystal
+end

src/SampledCorrelations/DataRetrieval.jl

@@ -44,37 +44,47 @@ end
     formula = intensity_formula(sc::SampledCorrelations; kwargs...)
     formula.calc_intensity(sc,q,ix_q,ix_ω)
 
-Establish a formula for computing the intensity of the discrete scattering modes `(q,ω)` using the correlation data ``𝒮^{αβ}(q,ω)`` stored in the [`SampledCorrelations`](@ref).
-The `formula` returned from `intensity_formula` can be passed to [`intensities_interpolated`](@ref) or [`intensities_binned`](@ref).
+Establish a formula for computing the intensity of the discrete scattering modes
+`(q,ω)` using the correlation data ``𝒮^{αβ}(q,ω)`` stored in the
+[`SampledCorrelations`](@ref). The `formula` returned from `intensity_formula`
+can be passed to [`intensities_interpolated`](@ref) or
+[`intensities_binned`](@ref).
 
-Sunny has several built-in formulas that can be selected by setting `contraction_mode` to one of these values:
+Sunny has several built-in formulas that can be selected by setting
+`contraction_mode` to one of these values:
 
-- `:perp` (default), which contracts ``𝒮^{αβ}(q,ω)`` with the dipole factor ``δ_{αβ} - q_{α}q_{β}``, returning the unpolarized intensity.
+- `:perp` (default), which contracts ``𝒮^{αβ}(q,ω)`` with the dipole factor
+  ``δ_{αβ} - q_{α}q_{β}``, returning the unpolarized intensity.
 - `:trace`, which yields ``\\operatorname{tr} 𝒮(q,ω) = ∑_α 𝒮^{αα}(q,ω)``
 - `:full`, which will return all elements ``𝒮^{αβ}(𝐪,ω)`` without contraction.
 
-Additionally, there are keyword arguments providing temperature and form factor corrections:
+Additionally, there are keyword arguments providing temperature and form factor
+corrections:
 
 - `kT`: If a temperature is provided, the intensities will be rescaled by a
     temperature- and ω-dependent classical-to-quantum factor. `kT` should be
     specified when making comparisons with spin wave calculations or
-    experimental data. If `kT` is not specified, infinite temperature (no correction) is assumed.
+    experimental data. If `kT` is not specified, infinite temperature (no
+    correction) is assumed.
 - `formfactors`: To apply form factor corrections, provide this keyword with a
-    vector of `FormFactor`s, one for each unique site in the unit cell. The form factors
-    will be symmetry propagated to all equivalent sites.
+    vector of `FormFactor`s, one for each unique site in the unit cell. The form
+    factors will be symmetry propagated to all equivalent sites.
 
-Alternatively, a custom formula can be specifed by providing a function `intensity = f(q,ω,correlations)` and specifying which correlations it requires:
+Alternatively, a custom formula can be specifed by providing a function
+`intensity = f(q,ω,correlations)` and specifying which correlations it requires:
 
     intensity_formula(f,sc::SampledCorrelations, required_correlations; kwargs...)
 
-The function is intended to be specified using `do` notation. For example, this custom formula sums the off-diagonal correlations:
+The function is intended to be specified using `do` notation. For example, this
+custom formula sums the off-diagonal correlations:
 
     required = [(:Sx,:Sy),(:Sy,:Sz),(:Sx,:Sz)]
     intensity_formula(sc,required,return_type = ComplexF64) do k, ω, off_diagonal_correlations
         sum(off_diagonal_correlations)
     end
 
-If your custom formula returns a type other than `Float64`, use the `return_type` keyword argument to flag this.
+If your custom formula returns a type other than `Float64`, use the
+`return_type` keyword argument to flag this.
 """
 function intensity_formula(f::Function,sc::SampledCorrelations,required_correlations; kwargs...)
     # SQTODO: This corr_ix may contain repeated correlations if the user does a silly
@@ -94,7 +104,7 @@ function intensity_formula(f::Function,sc::SampledCorrelations,corr_ix::Abstract
     NAtoms = Val(size(sc.data)[2])
     NCorr = Val(length(corr_ix))
 
-    ωs_sc = ωs(sc;negative_energies=true)
+    ωs_sc = ωs(sc;negative_energies=true)  # I think this default `true` may be a problem
 
     # Intensity is calculated at the discrete (ix_q,ix_ω) modes available to the system.
     # Additionally, for momentum transfers outside of the first BZ, the norm `k` of the
@@ -110,29 +120,18 @@ function intensity_formula(f::Function,sc::SampledCorrelations,corr_ix::Abstract
 end
 
 function classical_to_quantum(ω, kT::Float64)
-  if kT == Inf
-    return 1.0
-  end
-  if ω > 0
-    ω/(kT*(1 - exp(-ω/kT)))
-  elseif iszero(ω)
-    1.0
-  else
-    -ω*exp(ω/kT)/(kT*(1 - exp(ω/kT)))
-  end
-end
-
-function prepare_form_factors(sc, formfactors)
-    if isnothing(formfactors)
-        cryst = isnothing(sc.origin_crystal) ? sc.crystal : sc.origin_crystal 
-        class_indices = [findfirst(==(class_label), cryst.classes) for class_label in unique(cryst.classes)]
-        formfactors = [FormFactor{Sunny.EMPTY_FF}(; atom) for atom in class_indices]
+    if kT == Inf
+        return 1.0
+    end
+    if ω > 0
+        ω/(kT*(1 - exp(-ω/kT)))
+    elseif iszero(ω)
+        1.0
+    else
+        -ω*exp(ω/kT)/(kT*(1 - exp(ω/kT)))
     end
-    formfactors = upconvert_form_factors(formfactors) # Ensure formfactors have consistent type
-    return propagate_form_factors(sc, formfactors)
 end
 
-
 """
     lorentzian(x, η) 
 

src/SampledCorrelations/FormFactor.jl

@@ -15,6 +15,8 @@ end
 Basic type for specifying form factor parameters. Must be provided a site within
 the unit cell (`atom`) and a string specifying the element name. This used when
 creating an [`intensity_formula`](@ref), which accepts a list of `FormFactors`s.
+Note that these corrections assume S(q,ω) is being calculated in the dipole
+approximation; they may not be appropriate for a general SU(N) system.
 
 A list of supported element names is available at:
 
@@ -103,6 +105,8 @@ end
 
 # If necessary, update the indices of FormFactors from original crystal
 # to the corresponding indices of the new crystal.
+# TODO: Moving ion information into `SpinInfo` will eliminate the need for this as well
+# as symmetry propagation.
 function map_form_factors_to_crystal(sc::SampledCorrelations, ffs::Vector{FormFactor{FFType}}) where FFType
     (; crystal, origin_crystal) = sc
 
@@ -121,6 +125,8 @@ function map_form_factors_to_crystal(sc::SampledCorrelations, ffs::Vector{FormFa
     return ffs_new
 end
 
+# Remap form factors to reshaped crystal (if necessary) and propagate to symmetry
+# equivalent sites.
 function propagate_form_factors(sc::SampledCorrelations, ffs::Vector{FormFactor{FFType}}) where FFType
     ffs = map_form_factors_to_crystal(sc, ffs)
     ref_atoms = [ff.atom for ff in ffs]
@@ -133,7 +139,20 @@ function propagate_form_factors(sc::SampledCorrelations, ffs::Vector{FormFactor{
     end
 end
 
+# Process form factor information into optimal form for efficient calculation.
+# Form factors are provided at the level of an `intensities` call, so there is
+# no opportunity to precalculate this in current arrangement. 
+function prepare_form_factors(sc, formfactors)
+    if isnothing(formfactors)
+        cryst = isnothing(sc.origin_crystal) ? sc.crystal : sc.origin_crystal 
+        class_indices = [findfirst(==(class_label), cryst.classes) for class_label in unique(cryst.classes)]
+        formfactors = [FormFactor{Sunny.EMPTY_FF}(; atom) for atom in class_indices]
+    end
+    formfactors = upconvert_form_factors(formfactors) # Ensure formfactors have consistent type
+    return propagate_form_factors(sc, formfactors)
+end
 
+# Second-order form factor calculation.
 function compute_form(k::Float64, params::FormFactor{DOUBLE_FF})
     s = k/4π
     g = params.g_lande
@@ -147,12 +166,14 @@ function compute_form(k::Float64, params::FormFactor{DOUBLE_FF})
     return ((2-g)/g) * (form2*s^2) + form1
 end
 
+# First-order form factor calculation.
 function compute_form(k::Float64, params::FormFactor{SINGLE_FF})
     s = k/4π
     (A, a, B, b, C, c, D) = params.J0_params
     return A*exp(-a*s^2) + B*exp(-b*s^2) + C*exp(-c*s^2) + D
 end
 
+# No form factor correction.
 function compute_form(::Float64, ::FormFactor{EMPTY_FF})
     return 1.0
 end

src/SampledCorrelations/SampledCorrelations.jl

@@ -9,8 +9,8 @@ struct SampledCorrelations{N}
     # 𝒮^{αβ}(q,ω) data and metadata
     data           :: Array{ComplexF64, 7}   # Raw SF data for 1st BZ (numcorrelations × natoms × natoms × latsize × energy)
     crystal        :: Crystal                # Crystal for interpretation of q indices in `data`
-    origin_crystal :: Union{Nothing,Crystal} # Original user-specified crystal (if different from above)  # FIXME: Eliminate
-    Δω             :: Float64                # Energy step size (could make this a virtual property)
+    origin_crystal :: Union{Nothing,Crystal} # Original user-specified crystal (if different from above) -- needed for FormFactor accounting
+    Δω             :: Float64                # Energy step size (could make this a virtual property)  
 
     # Correlation info (αβ indices of 𝒮^{αβ}(q,ω))
     observables    :: Vector{LinearMap}  # Operators corresponding to observables
@@ -110,7 +110,7 @@ function all_observable_names(sc::SampledCorrelations)
 end
 
 """
-    dynamic_correlations(sys::System; Δt, nω, ωmax, 
+    dynamical_correlations(sys::System; Δt, nω, ωmax, 
         process_trajectory=:none, observables=nothing, correlations=nothing) 
 
 Creates a `SampledCorrelations` for calculating and storing ``𝒮(𝐪,ω)`` data.