From e71ea8cd03f2fc9011df3c71e8a506f771e15064 Mon Sep 17 00:00:00 2001 From: Federico Cerisola Date: Mon, 16 Sep 2024 14:34:56 +0100 Subject: [PATCH 1/8] Moved spin dynamics time step function out of user facing routing This will allow more flexibility for testing improvements and detect potential type instabilities. --- src/Dynamics.jl | 76 +++++++++++++++++++++++++++---------------------- 1 file changed, 42 insertions(+), 34 deletions(-) diff --git a/src/Dynamics.jl b/src/Dynamics.jl index badca59..2267e3c 100644 --- a/src/Dynamics.jl +++ b/src/Dynamics.jl @@ -59,40 +59,9 @@ function diffeqsolver( push!(Cω2, matrix[i].C*transpose(matrix[i].C)) push!(b, t -> invsqrtS0*matrix[i].C*[bfield[i][1](t), bfield[i][2](t), bfield[i][3](t)]); end - function f(du, u, (Cω2v, Beff), t) - Cω2v = get_tmp(Cω2v, u) - Beff = get_tmp(Beff, u) - s = @view u[1:3*N] - v = @view u[1+3*N:3*N+3*M] - w = @view u[1+3*N+3*M:3*N+6*M] - ds = @view du[1:3*N] - dv = @view du[1+3*N:3*N+3*M] - dw = @view du[1+3*N+3*M:3*N+6*M] - for i in 1:N - Beff[i, :] .= Bext - for j in 1:M - Beff[i, :] .+= bcoupling[j][i]*(b[j](t) + mul!(Cω2v, Cω2[j], v[1+(j-1)*3:3+(j-1)*3])) - end - end - for i in 1:N - ds[1+(i-1)*3] = -(s[2+(i-1)*3]*Beff[i,3]-s[3+(i-1)*3]*Beff[i,2]) - ds[2+(i-1)*3] = -(s[3+(i-1)*3]*Beff[i,1]-s[1+(i-1)*3]*Beff[i,3]) - ds[3+(i-1)*3] = -(s[1+(i-1)*3]*Beff[i,2]-s[2+(i-1)*3]*Beff[i,1]) - for j in 1:N - ds[1+(i-1)*3] += -(s[2+(i-1)*3]*JH[i,j]*s[3+(j-1)*3]-s[3+(i-1)*3]*JH[i,j]*s[2+(j-1)*3]) - ds[2+(i-1)*3] += -(s[3+(i-1)*3]*JH[i,j]*s[1+(j-1)*3]-s[1+(i-1)*3]*JH[i,j]*s[3+(j-1)*3]) - ds[3+(i-1)*3] += -(s[1+(i-1)*3]*JH[i,j]*s[2+(j-1)*3]-s[2+(i-1)*3]*JH[i,j]*s[1+(j-1)*3]) - end - end - dv .= w - for i in 1:M - dw[1+3*(i-1):3+3*(i-1)] = -(Jlist[i].ω0^2)*v[1+3*(i-1):3+3*(i-1)] -Jlist[i].Γ*w[1+3*(i-1):3+3*(i-1)] - for j in 1:N - dw[1+3*(i-1):3+3*(i-1)] += -Jlist[i].α*bcoupling[i][j]*s[1+3*(j-1):3+3*(j-1)] - end - end - end - prob = ODEProblem(f, u0, tspan, (dualcache(zeros(3)), dualcache(zeros(N,3)))) + + params = (N, M, Bext, JH, Jlist, Cω2, b, bcoupling, dualcache(zeros(3)), dualcache(zeros(N,3))) + prob = ODEProblem(_spin_time_step!, u0, tspan, params) condition(u, t, integrator) = true function affect!(integrator) # projection for n in 1:N @@ -111,6 +80,45 @@ function diffeqsolver( return sol end +function _spin_time_step!( + du, + u, + (N, M, Bext, JH, Jlist, Cω2, b, bcoupling, Cω2v, Beff), + t +) + Cω2v = get_tmp(Cω2v, u) + Beff = get_tmp(Beff, u) + s = @view u[1:3*N] + v = @view u[1+3*N:3*N+3*M] + w = @view u[1+3*N+3*M:3*N+6*M] + ds = @view du[1:3*N] + dv = @view du[1+3*N:3*N+3*M] + dw = @view du[1+3*N+3*M:3*N+6*M] + for i in 1:N + Beff[i, :] .= Bext + for j in 1:M + Beff[i, :] .+= bcoupling[j][i]*(b[j](t) + mul!(Cω2v, Cω2[j], v[1+(j-1)*3:3+(j-1)*3])) + end + end + for i in 1:N + ds[1+(i-1)*3] = -(s[2+(i-1)*3]*Beff[i,3]-s[3+(i-1)*3]*Beff[i,2]) + ds[2+(i-1)*3] = -(s[3+(i-1)*3]*Beff[i,1]-s[1+(i-1)*3]*Beff[i,3]) + ds[3+(i-1)*3] = -(s[1+(i-1)*3]*Beff[i,2]-s[2+(i-1)*3]*Beff[i,1]) + for j in 1:N + ds[1+(i-1)*3] += -(s[2+(i-1)*3]*JH[i,j]*s[3+(j-1)*3]-s[3+(i-1)*3]*JH[i,j]*s[2+(j-1)*3]) + ds[2+(i-1)*3] += -(s[3+(i-1)*3]*JH[i,j]*s[1+(j-1)*3]-s[1+(i-1)*3]*JH[i,j]*s[3+(j-1)*3]) + ds[3+(i-1)*3] += -(s[1+(i-1)*3]*JH[i,j]*s[2+(j-1)*3]-s[2+(i-1)*3]*JH[i,j]*s[1+(j-1)*3]) + end + end + dv .= w + for i in 1:M + dw[1+3*(i-1):3+3*(i-1)] = -(Jlist[i].ω0^2)*v[1+3*(i-1):3+3*(i-1)] -Jlist[i].Γ*w[1+3*(i-1):3+3*(i-1)] + for j in 1:N + dw[1+3*(i-1):3+3*(i-1)] += -Jlist[i].α*bcoupling[i][j]*s[1+3*(j-1):3+3*(j-1)] + end + end +end + """ function diffeqsolver(s0, tspan, J::LorentzianSD, Jshared::LorentzianSD, bfields, bfieldshared, matrix::Coupling; JH=zero(I), S0=1/2, Bext=[0, 0, 1], saveat=[], save_fields=false, projection=true, alg=Tsit5(), atol=1e-3, rtol=1e-3) From 791a7f2cd9117b5121be2d7759fd943914df2335 Mon Sep 17 00:00:00 2001 From: Federico Cerisola Date: Mon, 16 Sep 2024 14:40:03 +0100 Subject: [PATCH 2/8] Fixed parameter type instability warning from DifferentialEquations --- src/Dynamics.jl | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/Dynamics.jl b/src/Dynamics.jl index 2267e3c..37482f6 100644 --- a/src/Dynamics.jl +++ b/src/Dynamics.jl @@ -53,12 +53,8 @@ function diffeqsolver( M = length(Jlist) u0 = [s0; zeros(6*M)] invsqrtS0 = 1/sqrt(S0) - Cω2 = [] - b = [] - for i in 1:M - push!(Cω2, matrix[i].C*transpose(matrix[i].C)) - push!(b, t -> invsqrtS0*matrix[i].C*[bfield[i][1](t), bfield[i][2](t), bfield[i][3](t)]); - end + Cω2 = [matrix[i].C*transpose(matrix[i].C) for i in 1:M] + b = [t -> invsqrtS0*matrix[i].C*[bfield[i][1](t), bfield[i][2](t), bfield[i][3](t)] for i in 1:M] params = (N, M, Bext, JH, Jlist, Cω2, b, bcoupling, dualcache(zeros(3)), dualcache(zeros(N,3))) prob = ODEProblem(_spin_time_step!, u0, tspan, params) From 21802f70876e928a217bfd79c8466a7e7325e657 Mon Sep 17 00:00:00 2001 From: Federico Cerisola Date: Mon, 16 Sep 2024 14:55:04 +0100 Subject: [PATCH 3/8] Added linebreaks for better code readability --- src/Dynamics.jl | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/Dynamics.jl b/src/Dynamics.jl index 37482f6..9d80fb0 100644 --- a/src/Dynamics.jl +++ b/src/Dynamics.jl @@ -84,31 +84,39 @@ function _spin_time_step!( ) Cω2v = get_tmp(Cω2v, u) Beff = get_tmp(Beff, u) + s = @view u[1:3*N] v = @view u[1+3*N:3*N+3*M] w = @view u[1+3*N+3*M:3*N+6*M] ds = @view du[1:3*N] dv = @view du[1+3*N:3*N+3*M] dw = @view du[1+3*N+3*M:3*N+6*M] + for i in 1:N Beff[i, :] .= Bext + for j in 1:M Beff[i, :] .+= bcoupling[j][i]*(b[j](t) + mul!(Cω2v, Cω2[j], v[1+(j-1)*3:3+(j-1)*3])) end end + for i in 1:N ds[1+(i-1)*3] = -(s[2+(i-1)*3]*Beff[i,3]-s[3+(i-1)*3]*Beff[i,2]) ds[2+(i-1)*3] = -(s[3+(i-1)*3]*Beff[i,1]-s[1+(i-1)*3]*Beff[i,3]) ds[3+(i-1)*3] = -(s[1+(i-1)*3]*Beff[i,2]-s[2+(i-1)*3]*Beff[i,1]) + for j in 1:N ds[1+(i-1)*3] += -(s[2+(i-1)*3]*JH[i,j]*s[3+(j-1)*3]-s[3+(i-1)*3]*JH[i,j]*s[2+(j-1)*3]) ds[2+(i-1)*3] += -(s[3+(i-1)*3]*JH[i,j]*s[1+(j-1)*3]-s[1+(i-1)*3]*JH[i,j]*s[3+(j-1)*3]) ds[3+(i-1)*3] += -(s[1+(i-1)*3]*JH[i,j]*s[2+(j-1)*3]-s[2+(i-1)*3]*JH[i,j]*s[1+(j-1)*3]) end end + dv .= w + for i in 1:M dw[1+3*(i-1):3+3*(i-1)] = -(Jlist[i].ω0^2)*v[1+3*(i-1):3+3*(i-1)] -Jlist[i].Γ*w[1+3*(i-1):3+3*(i-1)] + for j in 1:N dw[1+3*(i-1):3+3*(i-1)] += -Jlist[i].α*bcoupling[i][j]*s[1+3*(j-1):3+3*(j-1)] end From 08c5bdc16c9ef5e0aee7a1ec6e6fe0af3404c447 Mon Sep 17 00:00:00 2001 From: Federico Cerisola Date: Mon, 16 Sep 2024 15:09:37 +0100 Subject: [PATCH 4/8] Reduce allocations in spin dynamics time stepping function --- src/Dynamics.jl | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/src/Dynamics.jl b/src/Dynamics.jl index 9d80fb0..a777e9b 100644 --- a/src/Dynamics.jl +++ b/src/Dynamics.jl @@ -94,9 +94,13 @@ function _spin_time_step!( for i in 1:N Beff[i, :] .= Bext + end + + for j in 1:M + vj = @view v[1+(j-1)*3:3+(j-1)*3] - for j in 1:M - Beff[i, :] .+= bcoupling[j][i]*(b[j](t) + mul!(Cω2v, Cω2[j], v[1+(j-1)*3:3+(j-1)*3])) + for i in 1:N + Beff[i, :] .+= bcoupling[j][i]*(b[j](t) + mul!(Cω2v, Cω2[j], vj)) end end @@ -115,10 +119,19 @@ function _spin_time_step!( dv .= w for i in 1:M - dw[1+3*(i-1):3+3*(i-1)] = -(Jlist[i].ω0^2)*v[1+3*(i-1):3+3*(i-1)] -Jlist[i].Γ*w[1+3*(i-1):3+3*(i-1)] + vi = @view v[1+(i-1)*3:3+(i-1)*3] + wi = @view w[1+(i-1)*3:3+(i-1)*3] + + for k in 1:3 + dw[k+3*(i-1)] = -(Jlist[i].ω0^2)*vi[k] -Jlist[i].Γ*wi[k] + end for j in 1:N - dw[1+3*(i-1):3+3*(i-1)] += -Jlist[i].α*bcoupling[i][j]*s[1+3*(j-1):3+3*(j-1)] + sj = @view s[1+3*(j-1):3+3*(j-1)] + + for k in 1:3 + dw[k+3*(i-1)] += -Jlist[i].α*bcoupling[i][j]*sj[k] + end end end end From 7fddae5e70044b499d7c4ad16bbbb20d184ce140 Mon Sep 17 00:00:00 2001 From: Federico Cerisola Date: Mon, 16 Sep 2024 15:27:50 +0100 Subject: [PATCH 5/8] Make spin dynamics time step function completely allocation free This required some internal change on how the bfields are handled. --- src/Dynamics.jl | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/Dynamics.jl b/src/Dynamics.jl index a777e9b..80d15a2 100644 --- a/src/Dynamics.jl +++ b/src/Dynamics.jl @@ -53,10 +53,11 @@ function diffeqsolver( M = length(Jlist) u0 = [s0; zeros(6*M)] invsqrtS0 = 1/sqrt(S0) + Cω = [matrix[i].C for i in 1:M] Cω2 = [matrix[i].C*transpose(matrix[i].C) for i in 1:M] - b = [t -> invsqrtS0*matrix[i].C*[bfield[i][1](t), bfield[i][2](t), bfield[i][3](t)] for i in 1:M] - params = (N, M, Bext, JH, Jlist, Cω2, b, bcoupling, dualcache(zeros(3)), dualcache(zeros(N,3))) + b, Cb, Cω2v, Beff = dualcache(zeros(3)), dualcache(zeros(3)), dualcache(zeros(3)), dualcache(zeros(N, 3)) + params = (N, M, invsqrtS0, Bext, JH, Jlist, Cω, Cω2, bfield, bcoupling, b, Cb, Cω2v, Beff) prob = ODEProblem(_spin_time_step!, u0, tspan, params) condition(u, t, integrator) = true function affect!(integrator) # projection @@ -79,9 +80,11 @@ end function _spin_time_step!( du, u, - (N, M, Bext, JH, Jlist, Cω2, b, bcoupling, Cω2v, Beff), + (N, M, invsqrtS0, Bext, JH, Jlist, Cω, Cω2, bfields, bcoupling, b, Cb, Cω2v, Beff), t ) + b = get_tmp(b, u) + Cb = get_tmp(Cb, u) Cω2v = get_tmp(Cω2v, u) Beff = get_tmp(Beff, u) @@ -98,9 +101,20 @@ function _spin_time_step!( for j in 1:M vj = @view v[1+(j-1)*3:3+(j-1)*3] + + for k in 1:3 + b[k] = bfields[j][k](t) + end + + mul!(Cb, Cω[j], b) + lmul!(invsqrtS0, Cb) + mul!(Cω2v, Cω2[j], vj) + Cb .+= Cω2v for i in 1:N - Beff[i, :] .+= bcoupling[j][i]*(b[j](t) + mul!(Cω2v, Cω2[j], vj)) + for k in 1:3 + Beff[i, k] += bcoupling[j][i]*Cb[k] + end end end From 0afc78c7cf3cf32f36226a8c9a56259fe1bee299 Mon Sep 17 00:00:00 2001 From: Federico Cerisola Date: Mon, 16 Sep 2024 15:46:44 +0100 Subject: [PATCH 6/8] Clean up code in spin dynamics time steping function Use more @views to make the code more readable without any performance cost. --- src/Dynamics.jl | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/src/Dynamics.jl b/src/Dynamics.jl index 80d15a2..bce96fa 100644 --- a/src/Dynamics.jl +++ b/src/Dynamics.jl @@ -119,14 +119,19 @@ function _spin_time_step!( end for i in 1:N - ds[1+(i-1)*3] = -(s[2+(i-1)*3]*Beff[i,3]-s[3+(i-1)*3]*Beff[i,2]) - ds[2+(i-1)*3] = -(s[3+(i-1)*3]*Beff[i,1]-s[1+(i-1)*3]*Beff[i,3]) - ds[3+(i-1)*3] = -(s[1+(i-1)*3]*Beff[i,2]-s[2+(i-1)*3]*Beff[i,1]) + si = @view s[1+(i-1)*3:3+(i-1)*3] + dsi = @view ds[1+(i-1)*3:3+(i-1)*3] + + dsi[1] = -(si[2]*Beff[i,3] - si[3]*Beff[i,2]) + dsi[2] = -(si[3]*Beff[i,1] - si[1]*Beff[i,3]) + dsi[3] = -(si[1]*Beff[i,2] - si[2]*Beff[i,1]) for j in 1:N - ds[1+(i-1)*3] += -(s[2+(i-1)*3]*JH[i,j]*s[3+(j-1)*3]-s[3+(i-1)*3]*JH[i,j]*s[2+(j-1)*3]) - ds[2+(i-1)*3] += -(s[3+(i-1)*3]*JH[i,j]*s[1+(j-1)*3]-s[1+(i-1)*3]*JH[i,j]*s[3+(j-1)*3]) - ds[3+(i-1)*3] += -(s[1+(i-1)*3]*JH[i,j]*s[2+(j-1)*3]-s[2+(i-1)*3]*JH[i,j]*s[1+(j-1)*3]) + sj = @view s[1+(j-1)*3:3+(j-1)*3] + + dsi[1] += -(si[2]*JH[i,j]*sj[3] - si[3]*JH[i,j]*sj[2]) + dsi[2] += -(si[3]*JH[i,j]*sj[1] - si[1]*JH[i,j]*sj[3]) + dsi[3] += -(si[1]*JH[i,j]*sj[2] - si[2]*JH[i,j]*sj[1]) end end @@ -135,16 +140,17 @@ function _spin_time_step!( for i in 1:M vi = @view v[1+(i-1)*3:3+(i-1)*3] wi = @view w[1+(i-1)*3:3+(i-1)*3] + dwi = @view dw[1+(i-1)*3:3+(i-1)*3] for k in 1:3 - dw[k+3*(i-1)] = -(Jlist[i].ω0^2)*vi[k] -Jlist[i].Γ*wi[k] + dwi[k] = -(Jlist[i].ω0^2)*vi[k] -Jlist[i].Γ*wi[k] end for j in 1:N sj = @view s[1+3*(j-1):3+3*(j-1)] for k in 1:3 - dw[k+3*(i-1)] += -Jlist[i].α*bcoupling[i][j]*sj[k] + dwi[k] += -Jlist[i].α*bcoupling[i][j]*sj[k] end end end From 1511be296e3b1fd51779114a810e84ba41f3afe0 Mon Sep 17 00:00:00 2001 From: Federico Cerisola Date: Mon, 16 Sep 2024 15:49:36 +0100 Subject: [PATCH 7/8] Add whitespace for better readability --- src/Dynamics.jl | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/Dynamics.jl b/src/Dynamics.jl index bce96fa..69014f8 100644 --- a/src/Dynamics.jl +++ b/src/Dynamics.jl @@ -50,6 +50,7 @@ function diffeqsolver( if length(Jlist) != length(matrix) || length(Jlist) != length(bcoupling) throw(DimensionMismatch("The dimension of Jlist, bcoupling, and matrix must match.")) end + M = length(Jlist) u0 = [s0; zeros(6*M)] invsqrtS0 = 1/sqrt(S0) @@ -59,6 +60,7 @@ function diffeqsolver( b, Cb, Cω2v, Beff = dualcache(zeros(3)), dualcache(zeros(3)), dualcache(zeros(3)), dualcache(zeros(N, 3)) params = (N, M, invsqrtS0, Bext, JH, Jlist, Cω, Cω2, bfield, bcoupling, b, Cb, Cω2v, Beff) prob = ODEProblem(_spin_time_step!, u0, tspan, params) + condition(u, t, integrator) = true function affect!(integrator) # projection for n in 1:N @@ -68,12 +70,15 @@ function diffeqsolver( end cb = DiscreteCallback(condition, affect!, save_positions=(false,false)) skwargs = projection ? (callback=cb,) : NamedTuple() + if save_fields save_idxs = 1:(3*N+6*M) else save_idxs = 1:3*N end + sol = solve(prob, alg; abstol=atol, reltol=rtol, maxiters=Int(1e9), save_idxs=save_idxs, saveat=saveat, kwargs..., skwargs...) + return sol end From d4459a2796e81e291cefdccd1721fa0ed4734abe Mon Sep 17 00:00:00 2001 From: Federico Cerisola Date: Mon, 16 Sep 2024 16:12:17 +0100 Subject: [PATCH 8/8] Remove explicit loops over vector components Makes code more readable without performance penalty --- src/Dynamics.jl | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/src/Dynamics.jl b/src/Dynamics.jl index 69014f8..2cb20ec 100644 --- a/src/Dynamics.jl +++ b/src/Dynamics.jl @@ -117,9 +117,9 @@ function _spin_time_step!( Cb .+= Cω2v for i in 1:N - for k in 1:3 - Beff[i, k] += bcoupling[j][i]*Cb[k] - end + Beffi = @view Beff[i, :] + + @. Beffi += bcoupling[j][i]*Cb end end @@ -147,16 +147,12 @@ function _spin_time_step!( wi = @view w[1+(i-1)*3:3+(i-1)*3] dwi = @view dw[1+(i-1)*3:3+(i-1)*3] - for k in 1:3 - dwi[k] = -(Jlist[i].ω0^2)*vi[k] -Jlist[i].Γ*wi[k] - end + @. dwi = -(Jlist[i].ω0^2)*vi - Jlist[i].Γ*wi for j in 1:N sj = @view s[1+3*(j-1):3+3*(j-1)] - for k in 1:3 - dwi[k] += -Jlist[i].α*bcoupling[i][j]*sj[k] - end + @. dwi += -Jlist[i].α*bcoupling[i][j]*sj end end end