Functional optim via ortho-stereo projection

SunnySuite · Jul 18, 2023 · 6859c73 · 6859c73
1 parent fe09d74
commit 6859c73
Show file tree

Hide file tree

Showing 2 changed files with 88 additions and 26 deletions.
diff --git a/src/Integrators.jl b/src/Integrators.jl
@@ -163,6 +163,31 @@ end
     return :(CVec{$N}($(out...)))
 end
 
+# Analog of set_forces! for Schroedinger dynamics. Placed here because will need
+# to be revised in tandem with rhs_langevin! and rhs_ll!. This could possibly be
+# factored out of those functions to reduce code duplication. This is performance
+# critical, so would need benchmarking.
+function set_complex_forces!(HZ, Z, sys::System{N}) where N
+    B = only(get_dipole_buffers(sys, 1) )
+
+    @. sys.dipoles = expected_spin(Z) # temporarily desyncs dipoles and coherents
+    set_forces!(B, sys.dipoles, sys)
+
+    if is_homogeneous(sys)
+        ints = interactions_homog(sys)
+        for site in all_sites(sys)
+            Λ = ints[to_atom(site)].aniso.matrep
+            HZ[site] = mul_spin_matrices(Λ, -B[site], Z[site])  
+        end 
+    else
+        ints = interactions_inhomog(sys)
+        for site in all_sites(sys)
+            Λ = ints[site].aniso.matrep
+            HZ[site] = mul_spin_matrices(Λ, -B[site], Z[site])  
+        end 
+    end
+end
+
 
 function rhs_langevin!(ΔZ::Array{CVec{N}, 4}, Z::Array{CVec{N}, 4}, ξ::Array{CVec{N}, 4},
                         B::Array{Vec3, 4}, integrator::Langevin, sys::System{N}) where N

diff --git a/src/Optimization.jl b/src/Optimization.jl
@@ -173,32 +173,6 @@ function minimize_energy_stereo!(sys; method=Optim.LBFGS, maxiters = 100, kwargs
     return success
 end
 
-################################################################################
-# Equivalent of set forces for SU(N) systems
-################################################################################
-
-# Calculates ℌz 
-function set_complex_forces!(HZ, Z, sys::System{N}) where N
-    B = only(get_dipole_buffers(sys, 1) )
-
-    @. sys.dipoles = expected_spin(Z) # temporarily desyncs dipoles and coherents
-    set_forces!(B, sys.dipoles, sys)
-
-    if is_homogeneous(sys)
-        ints = interactions_homog(sys)
-        for site in all_sites(sys)
-            Λ = ints[to_atom(site)].aniso.matrep
-            HZ[site] = mul_spin_matrices(Λ, -B[site], -Z[site])  # Z negative to match sign convention of set_forces!
-        end 
-    else
-        ints = interactions_inhomog(sys)
-        for site in all_sites(sys)
-            Λ = ints[site].aniso.matrep
-            HZ[site] = mul_spin_matrices(Λ, -B[site], -Z[site])  # Z negative to match sign convention of set_forces!
-        end 
-    end
-end
-
 ################################################################################
 # Coordinates for ℂP^{N-1} 
 ################################################################################
@@ -383,3 +357,66 @@ function there_and_back2(u::SVector{N, ComplexF64}) where N
     display(spin_expectations(u1))
     u1 * exp(-im*angle(u1[1]))
 end
+
+
+################################################################################
+# Kipton's projection
+################################################################################
+
+function projective_to_conventional(α::SVector{N, ComplexF64}, z::SVector{N, ComplexF64}) where N
+    v = (I - z*z')*α
+    v2 = v'*v
+    normalize((2v + (1-v2)*z) / (1+v2)) # Normalize shouldn't be necessary
+end
+
+# Make non-allocating version
+function projective_jac!(jac, α::SVector{N, ComplexF64}, z::SVector{N, ComplexF64}) where N
+
+    P = (I - z*z')
+    v = P*α
+    v2 = v' * v 
+    dv_dα = P 
+    dv2_dα = 2*(α' - 2*(α'*z)*z')
+
+    jac .= (1/(1+v2)) * (2dv_dα - z * dv2_dα)' - (1/(1+v2))^2 * dv2_dα' * ((2v + (1-v2)*z)')
+end
+
+function optim_energy_projective(αs, zs, sys::System{N}) where N
+    αs = reinterpret(reshape, SVector{N, ComplexF64}, αs)
+    for site in all_sites(sys)
+        set_coherent_state!(sys, projective_to_conventional(αs[site], zs[site]), site)
+    end
+    return energy(sys)
+end
+
+function optim_gradient_projective!(buf, αs, zs, sys::System{N}) where N
+    αs = reinterpret(reshape, SVector{N, ComplexF64}, αs)
+    Hgrad = reinterpret(reshape, SVector{N, ComplexF64}, buf)
+
+    # Calculate gradient of energy in original coordinates
+    for site in all_sites(sys)
+        set_coherent_state!(sys, projective_to_conventional(αs[site], zs[site]), site)
+    end
+    Sunny.set_complex_forces!(Hgrad, sys.coherents, sys)
+
+    # Calculate gradient, using Jacobian for projectivegraphic coordinates 
+    jac = zeros(ComplexF64, N, N)  # Maybe write function to apply matrix multiply and avoid allocation
+    for site in all_sites(sys)
+        projective_jac!(jac, αs[site], zs[site])
+        Hgrad[site] = jac * Hgrad[site]  # Note Optim expects ∇, `set_forces!` gives -∇
+    end
+end
+
+function minimize_energy_projective!(sys::System{N}; method=Optim.LBFGS, maxiters = 1000, kwargs...) where N
+
+    zs = copy(sys.coherents)
+    # αs = zeros(SVector{N, ComplexF64}, size(sys.coherents))
+    # αs = Array(reinterpret(reshape, ComplexF64, αs))
+    αs = zeros(ComplexF64, (N, size(sys.coherents)...))
+
+    f(proj_coords) = optim_energy_projective(proj_coords, zs, sys)
+    g!(G, proj_coords) = optim_gradient_projective!(G, proj_coords, zs, sys)
+
+    options = Optim.Options(iterations=maxiters, kwargs...)
+    Optim.optimize(f, g!, αs, method(), options)
+end