Optimization

The `foreachbond` abstraction was actually slowing various tests down.

src/Optimization.jl

@@ -99,7 +99,7 @@ function minimize_energy!(sys::System{N}; maxiters=100, subiters=10, method=Opti
     # Functions to calculate energy and gradient for the state `αs`
     function f(αs)
         optim_set_spins!(sys, αs, ns)
-        return energy(sys) # TODO: `Sunny.energy` seems to allocate and is type-unstable (7/20/2023)
+        return energy(sys)
     end
     function g!(G, αs)
         optim_set_spins!(sys, αs, ns)

src/Sunny.jl

@@ -33,8 +33,6 @@ const HermitianC64 = Hermitian{ComplexF64, Matrix{ComplexF64}}
 
 @static if VERSION < v"1.10" hermitianpart(A) = Hermitian(A+A')/2 end
 
-include("Util/CartesianIndicesShifted.jl")
-
 include("Operators/Spin.jl")
 include("Operators/Rotation.jl")
 include("Operators/Stevens.jl")

src/System/Interactions.jl

@@ -203,130 +203,130 @@ end
 The total system energy. See also [`energy_per_site`](@ref).
 """
 function energy(sys::System{N}) where N
-    (; crystal, latsize, dipoles, extfield, ewald) = sys
-
     E = 0.0
 
     # Zeeman coupling to external field
     for site in eachsite(sys)
-        E -= sys.units.μB * extfield[site] ⋅ (sys.gs[site] * dipoles[site])
+        E -= sys.units.μB * sys.extfield[site] ⋅ (sys.gs[site] * sys.dipoles[site])
     end
 
     # Anisotropies and exchange interactions
-    for i in 1:natoms(crystal)
+    for i in 1:natoms(sys.crystal)
         if is_homogeneous(sys)
+            # Interactions for sublattice i (same for every cell)
             interactions = sys.interactions_union[i]
-            E += energy_aux(sys, interactions, i, eachcell(sys), homog_bond_iterator(latsize))
+            E += energy_aux(interactions, sys, i, eachcell(sys))
         else
             for cell in eachcell(sys)
                 interactions = sys.interactions_union[cell, i]
-                E += energy_aux(sys, interactions, i, (cell,), inhomog_bond_iterator(latsize, cell))
+                E += energy_aux(interactions, sys, i, (cell,))
             end
         end
     end
 
     # Long-range dipole-dipole
-    if !isnothing(ewald)
+    if !isnothing(sys.ewald)
         E += ewald_energy(sys)
     end
 
     return E
 end
 
 # Total energy contributed by sublattice `i`, summed over the list of `cells`.
-# The function `foreachbond` enables efficient iteration over neighboring cell
-# pairs (without double counting).
-function energy_aux(sys::System{N}, ints::Interactions, i::Int, cells, foreachbond) where N
-    (; dipoles, coherents) = sys
+function energy_aux(ints::Interactions, sys::System{N}, i::Int, cells) where N
     E = 0.0
 
     # Single-ion anisotropy
     if N == 0       # Dipole mode
         stvexp = ints.onsite :: StevensExpansion
         for cell in cells
-            s = dipoles[cell, i]
+            s = sys.dipoles[cell, i]
             E += energy_and_gradient_for_classical_anisotropy(s, stvexp)[1]
         end
     else            # SU(N) mode
         Λ = ints.onsite :: HermitianC64
         for cell in cells
-            Z = coherents[cell, i]
+            Z = sys.coherents[cell, i]
             E += real(dot(Z, Λ, Z))
         end
     end
 
-    foreachbond(ints.pair) do pc, site1, site2
-        sᵢ = dipoles[site1]
-        sⱼ = dipoles[site2]
-
-        # Bilinear
-        J = pc.bilin
-        E += dot(sᵢ, J, sⱼ)
-
-        # Biquadratic
-        if !iszero(pc.biquad)
-            J = pc.biquad
-            if sys.mode == :dipole
-                E += J * (sᵢ⋅sⱼ)^2
-            elseif sys.mode == :SUN
-                error("Biquadratic currently unsupported in SU(N) mode.")
+    for pc in ints.pair
+        (; bond, isculled) = pc
+        isculled && break
+
+        for cellᵢ in cells
+            cellⱼ = offsetc(cellᵢ, bond.n, sys.latsize)
+            sᵢ = sys.dipoles[cellᵢ, bond.i]
+            sⱼ = sys.dipoles[cellⱼ, bond.j]
+
+            # Bilinear
+            J = pc.bilin :: Union{Float64, Mat3}
+            E += dot(sᵢ, J, sⱼ)
+
+            # Biquadratic
+            if !iszero(pc.biquad)
+                J = pc.biquad
+                if sys.mode == :dipole
+                    E += J * (sᵢ⋅sⱼ)^2
+                elseif sys.mode == :SUN
+                    error("Biquadratic currently unsupported in SU(N) mode.")
+                end
             end
-        end
 
-        # General
-        if N > 0
-            Zᵢ = coherents[site1]
-            Zⱼ = coherents[site2]
-            for (A, B) in pc.general.data
-                Ā = real(dot(Zᵢ, A, Zᵢ))
-                B̄ = real(dot(Zⱼ, B, Zⱼ))
-                E += Ā * B̄
+            # General
+            if N > 0
+                Zᵢ = sys.coherents[cellᵢ, bond.i]
+                Zⱼ = sys.coherents[cellⱼ, bond.j]
+                for (A, B) in pc.general.data
+                    Ā = real(dot(Zᵢ, A, Zᵢ))
+                    B̄ = real(dot(Zⱼ, B, Zⱼ))
+                    E += Ā * B̄
+                end
             end
         end
     end
 
     return E
 end
 
+
 # Updates ∇E in-place to hold energy gradient, dE/ds, for each spin. In the case
 # of :SUN mode, s is interpreted as expected spin, and dE/ds only includes
 # contributions from Zeeman coupling, bilinear exchange, and long-range
 # dipole-dipole. Excluded terms include onsite coupling, and general pair
 # coupling (biquadratic and beyond).
 function set_energy_grad_dipoles!(∇E, dipoles::Array{Vec3, 4}, sys::System{N}) where N
-    (; crystal, latsize, extfield, ewald) = sys
-
     fill!(∇E, zero(Vec3))
 
     # Zeeman coupling
     for site in eachsite(sys)
-        ∇E[site] -= sys.units.μB * (sys.gs[site]' * extfield[site])
+        ∇E[site] -= sys.units.μB * (sys.gs[site]' * sys.extfield[site])
     end
 
     # Anisotropies and exchange interactions
-    for i in 1:natoms(crystal)
+    for i in 1:natoms(sys.crystal)
         if is_homogeneous(sys)
-            # Interaction is the same at every cell
+            # Interactions for sublattice i (same for every cell)
             interactions = sys.interactions_union[i]
-            set_energy_grad_dipoles_aux!(∇E, dipoles, interactions, sys, i, eachcell(sys), homog_bond_iterator(latsize))
+            set_energy_grad_dipoles_aux!(∇E, dipoles, interactions, sys, i, eachcell(sys))
         else
             for cell in eachcell(sys)
-                # There is a different interaction at every cell
-                interactions = sys.interactions_union[cell,i]
-                set_energy_grad_dipoles_aux!(∇E, dipoles, interactions, sys, i, (cell,), inhomog_bond_iterator(latsize, cell))
+                # Interactions for sublattice i and a specific cell
+                interactions = sys.interactions_union[cell, i]
+                set_energy_grad_dipoles_aux!(∇E, dipoles, interactions, sys, i, (cell,))
             end
         end
     end
 
-    if !isnothing(ewald)
+    if !isnothing(sys.ewald)
         accum_ewald_grad!(∇E, dipoles, sys)
     end
 end
 
 # Calculate the energy gradient `∇E' for the sublattice `i' at all elements of
-# `cells`. The function `foreachbond` enables efficient iteration over
-# neighboring cell pairs (without double counting).
-function set_energy_grad_dipoles_aux!(∇E, dipoles::Array{Vec3, 4}, ints::Interactions, sys::System{N}, i::Int, cells, foreachbond) where N
+# `cells`.
+function set_energy_grad_dipoles_aux!(∇E, dipoles::Array{Vec3, 4}, ints::Interactions, sys::System{N}, i::Int, cells) where N
     # Single-ion anisotropy only contributes in dipole mode. In SU(N) mode, the
     # anisotropy matrix will be incorporated directly into ℌ.
     if N == 0
@@ -337,23 +337,29 @@ function set_energy_grad_dipoles_aux!(∇E, dipoles::Array{Vec3, 4}, ints::Inter
         end
     end
 
-    foreachbond(ints.pair) do pc, site1, site2
-        sᵢ = dipoles[site1]
-        sⱼ = dipoles[site2]
-
-        # Bilinear
-        J = pc.bilin
-        ∇E[site1] += J  * sⱼ
-        ∇E[site2] += J' * sᵢ
-
-        # Biquadratic
-        if !iszero(pc.biquad)
-            J = pc.biquad
-            if sys.mode == :dipole
-                ∇E[site1] += J * 2sⱼ*(sᵢ⋅sⱼ)
-                ∇E[site2] += J * 2sᵢ*(sᵢ⋅sⱼ)
-            elseif sys.mode == :SUN
-                error("Biquadratic currently unsupported in SU(N) mode.")
+    for pc in ints.pair
+        (; bond, isculled) = pc
+        isculled && break
+
+        for cellᵢ in cells
+            cellⱼ = offsetc(cellᵢ, bond.n, sys.latsize)
+            sᵢ = dipoles[cellᵢ, bond.i]
+            sⱼ = dipoles[cellⱼ, bond.j]
+
+            # Bilinear
+            J = pc.bilin
+            ∇E[cellᵢ, bond.i] += J  * sⱼ
+            ∇E[cellⱼ, bond.j] += J' * sᵢ
+
+            # Biquadratic
+            if !iszero(pc.biquad)
+                J = pc.biquad
+                if sys.mode == :dipole
+                    ∇E[cellᵢ, bond.i] += J * 2sⱼ*(sᵢ⋅sⱼ)
+                    ∇E[cellⱼ, bond.j] += J * 2sᵢ*(sᵢ⋅sⱼ)
+                elseif sys.mode == :SUN
+                    error("Biquadratic currently unsupported in SU(N) mode.")
+                end
             end
         end
     end
@@ -421,34 +427,3 @@ end
     end
     return :(CVec{$N}($(out...)))
 end
-
-
-# Produces a function that iterates over a list interactions for a given cell
-function inhomog_bond_iterator(latsize, cell)
-    return function foreachbond(f, pcs)
-        for pc in pcs
-            # Early return to avoid double-counting a bond
-            pc.isculled && break
-
-            # Neighboring cell may wrap the system
-            cell′ = offsetc(cell, pc.bond.n, latsize)
-            f(pc, CartesianIndex(cell, pc.bond.i), CartesianIndex(cell′, pc.bond.j))
-        end
-    end
-end
-
-# Produces a function that iterates over a list of interactions, involving all
-# pairs of cells in a homogeneous system
-function homog_bond_iterator(latsize)
-    return function foreachbond(f, pcs)
-        for pc in pcs
-            # Early return to avoid double-counting a bond
-            pc.isculled && break
-
-            # Iterate over all cells and periodically shifted neighbors
-            for (ci, cj) in zip(CartesianIndices(latsize), CartesianIndicesShifted(latsize, Tuple(pc.bond.n)))
-                f(pc, CartesianIndex(ci, pc.bond.i), CartesianIndex(cj, pc.bond.j))
-            end
-        end
-    end
-end

test/test_jet.jl

@@ -9,20 +9,18 @@
         set_exchange!(sys, -1.0, Bond(1,1,(1,0,0)))
         polarize_spins!(sys, (0,0,1))
 
-        # Test type stability of LocalSampler
+        @test_opt energy(sys)
+
         sampler = LocalSampler(kT=0.2, propose=propose_flip)
         @test_opt step!(sys, sampler)
 
-        # Test type stability with mixed proposals
         propose = @mix_proposals 0.5 propose_flip 0.5 propose_delta(0.2)
         sampler = LocalSampler(kT=0.2; propose)
         @test_opt step!(sys, sampler)
 
-        # Test type stability of Langevin
         langevin = Langevin(0.01, kT=0.2, λ=0.1)
         @test_opt step!(sys, langevin)
 
-        # Test type stability of ImplicitMidpoint
         integrator = ImplicitMidpoint(0.01)
         @test_opt step!(sys, integrator)
     end
@@ -42,6 +40,11 @@ end
         set_exchange!(sys, -1.0, Bond(1,1,(1,0,0)))
         polarize_spins!(sys, (0,0,1))
 
+        # TODO: Diagonose possible @allocated bug. BenchmarkTools.@btime
+        # suggests that this is actually zero-allocation.
+        energy(sys)
+        @test 16 >= @allocated energy(sys)
+
         propose = @mix_proposals 0.5 propose_flip 0.5 propose_delta(0.2)
         sampler = LocalSampler(kT=0.2; propose)
         step!(sys, sampler)