From c19dfbec28350d8c8fb4e4edf81460cfe49069ca Mon Sep 17 00:00:00 2001
From: Kipton Barros <kbarros@gmail.com>
Date: Sat, 30 Sep 2023 08:15:29 -0600
Subject: [PATCH] Make `energy` type stable

---
 src/Optimization.jl        |  2 +-
 src/System/Interactions.jl | 16 ++++++++--------
 test/test_jet.jl           |  6 ++----
 3 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/src/Optimization.jl b/src/Optimization.jl
index dd6abca63..1d10de262 100644
--- a/src/Optimization.jl
+++ b/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)
diff --git a/src/System/Interactions.jl b/src/System/Interactions.jl
index f0fe459a7..9c527c9b0 100644
--- a/src/System/Interactions.jl
+++ b/src/System/Interactions.jl
@@ -238,20 +238,20 @@ end
 # pairs (without double counting).
 function energy_aux(sys::System{N}, ints::Interactions, i::Int, cells, foreachbond) where N
     (; dipoles, coherents) = sys
-    E = 0.0
+    E = Ref(0.0) # Type stable accumulator to be captured by closure.
 
     # Single-ion anisotropy
     if N == 0       # Dipole mode
         stvexp = ints.onsite :: StevensExpansion
         for cell in cells
             s = dipoles[cell, i]
-            E += energy_and_gradient_for_classical_anisotropy(s, stvexp)[1]
+            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]
-            E += real(dot(Z, Λ, Z))
+            E[] += real(dot(Z, Λ, Z))
         end
     end
 
@@ -260,14 +260,14 @@ function energy_aux(sys::System{N}, ints::Interactions, i::Int, cells, foreachbo
         sⱼ = dipoles[site2]
 
         # Bilinear
-        J = pc.bilin
-        E += dot(sᵢ, J, sⱼ)
+        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
+                E[] += J * (sᵢ⋅sⱼ)^2
             elseif sys.mode == :SUN
                 error("Biquadratic currently unsupported in SU(N) mode.")
             end
@@ -280,12 +280,12 @@ function energy_aux(sys::System{N}, ints::Interactions, i::Int, cells, foreachbo
             for (A, B) in pc.general.data
                 Ā = real(dot(Zᵢ, A, Zᵢ))
                 B̄ = real(dot(Zⱼ, B, Zⱼ))
-                E += Ā * B̄
+                E[] += Ā * B̄
             end
         end
     end
 
-    return E
+    return E[]
 end
 
 # Updates ∇E in-place to hold energy gradient, dE/ds, for each spin. In the case
diff --git a/test/test_jet.jl b/test/test_jet.jl
index 6f4f66679..2024fc18f 100644
--- a/test/test_jet.jl
+++ b/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