adaptive kink rho propagation

Author: amartyabose

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumDynamics"
 uuid = "5a40c529-53c2-4483-a223-e00c1cee8134"
 authors = ["Amartya Bose <[email protected]> and contributors"]
-version = "0.2.22"
+version = "0.2.23"
 
 [deps]
 AtomsIO = "1692102d-eeb4-4df9-807b-c9517f998d44"

src/DynamicMap_MasterEquation/TTM.jl

@@ -172,7 +172,7 @@ function get_propagators(; fbU::Array{<:Complex,3}, Jw::Vector{T}, β, dt, ntime
         end
         T0e = get_Ts(U0e_within_r)
         if !isnothing(output)
-            Utilities.check_or_insert_value(output, "Tmat", T0e)
+            Utilities.check_or_insert_value(output, "T0e", T0e)
             flush(output)
         end
 

src/PathIntegral/QuAPI.jl

@@ -570,4 +570,77 @@ function build_augmented_propagator_kink(; fbU::AbstractArray{ComplexF64,3}, Jw:
     end
     U0e
 end
+
+function propagate_kink(; fbU::AbstractArray{ComplexF64,3}, Jw::Vector{T}, ρ0::AbstractMatrix{ComplexF64}, β::Real, dt::Real, ntimes::Int, extraargs::QuAPIArgs=QuAPIArgs(), svec=[1.0 -1.0], reference_prop=false, verbose::Bool=false, output::Union{Nothing,HDF5.Group}=nothing, from_TTM::Bool=false, exec=ThreadedEx()) where {T<:SpectralDensities.SpectralDensity}
+    @assert length(Jw) == size(svec, 1)
+    η = [EtaCoefficients.calculate_η(jw; β, dt, kmax=ntimes, imaginary_only=reference_prop) for jw in Jw]
+    sdim = size(fbU, 2)
+    sdim2 = sdim^2
+    state_values, _ = setup_simulation(ones(sdim, sdim), η, svec, extraargs)
+
+    if verbose
+        @info "Starting propagation within memory"
+    end
+    ρs = Array{ComplexF64}(undef, ntimes+1, sdim, sdim)
+    ρs[1, :, :] = ρ0
+    if !isnothing(output)
+        Utilities.check_or_insert_value(output, "rho", ρs)
+        Utilities.check_or_insert_value(output, "time_taken", zeros(Float64, ntimes))
+        Utilities.check_or_insert_value(output, "num_paths", zeros(Int64, ntimes))
+    end
+    nkinks = extraargs.num_kinks==-1 ? ntimes : extraargs.num_kinks
+    nblips = extraargs.num_blips==-1 ? ntimes+1 : extraargs.num_blips
+    numpaths = zeros(Int64, ntimes)
+    timetaken = zeros(Float64, ntimes)
+    for ind in findall(!iszero, ρ0)
+        forward_paths = [[UInt8(ind[1])]]
+        forward_amplitudes = [1.0+0.0im]
+        backward_paths = [[UInt8(ind[2])]]
+        backward_amplitudes = [1.0+0.0im]
+        for i = 1:ntimes
+            if verbose
+                @info "Starting time step $(i)"
+            end
+            _, time_taken, memory_allocated, gc_time, _ = @timed begin
+                forward_paths, forward_amplitudes = Utilities.generate_paths_kink_limit(forward_paths, forward_amplitudes, nkinks, sdim, fbU, extraargs.prop_cutoff, sqrt(extraargs.cutoff))
+                backward_paths, backward_amplitudes = Utilities.generate_paths_kink_limit(backward_paths, backward_amplitudes, nkinks, sdim, fbU, extraargs.prop_cutoff, sqrt(extraargs.cutoff))
+                if verbose
+                    @info "Path generation done. $(length(forward_paths)) forward paths generated. Expect up to $(length(forward_paths)^2) forward-backward paths based on filtration."
+                end
+                @floop exec for ((fp, fa), (bp, ba)) in Iterators.product(zip(forward_paths, forward_amplitudes), zip(backward_paths, backward_amplitudes))
+                    num_blips = count(fp .!= bp)
+                    if num_blips > nblips
+                        continue
+                    end
+                    bare_amplitude = fa * conj(ba)
+                    if abs(bare_amplitude) < extraargs.cutoff
+                        continue
+                    end
+                    @init states = zeros(UInt16, i+1)
+                    states .= (fp.-1) .* sdim .+ bp
+                    @reduce num_paths = 0 + 1
+                    for (bn, bη) in enumerate(η)
+                        @inbounds bare_amplitude *= get_path_influence(bη, bn, state_values, states, false)
+                    end
+                    @init tmprho = zeros(ComplexF64, sdim2)
+                    tmprho .= 0.0
+                    @inbounds tmprho[states[end]] = bare_amplitude * ρ0[ind]
+                    @reduce tmprhot = zeros(ComplexF64, sdim2) .+ tmprho
+                end
+                numpaths[i] += num_paths
+                @inbounds ρs[i+1, :, :] .+= Utilities.density_matrix_vector_to_matrix(tmprhot)
+            end
+            timetaken[i] += time_taken
+            if verbose
+                @info "Done time step $(i); # paths = $(sum(num_paths)); time = $(round(time_taken; digits=3)) sec; memory allocated = $(round(memory_allocated / 1e6; digits=3)) GB; gc time = $(round(gc_time; digits=3)) sec"
+            end
+            if !isnothing(output)
+                output["rho"][i+1, :, :] = ρs
+                output["time_taken"][i] = timetaken[i]
+                output["num_paths"] = numpaths
+            end
+        end
+    end
+    0:dt:ntimes*dt, ρs
+end
 end

test/runtests.jl

@@ -72,6 +72,16 @@ end
         @test all(norm.(ρs[:, 1, 2] .- complex_dat[1:ntimes+1, 2, 1]) .< 1e-2)
     end
 
+    @testset "adaptive-kink-QuAPI" begin
+        fU = Propagators.calculate_bare_propagators(; Hamiltonian, dt, ntimes, forward_backward=false)
+        t, ρs = QuAPI.propagate_kink(; fbU=fU, Jw=[Jw], β, ρ0, dt, ntimes, svec)
+
+        @test all(ρs[:, 1, 2] .≈ conj(ρs[:, 2, 1]))
+        @test all(ρs[:, 1, 1] .+ ρs[:, 2, 2] .≈ 1.0 + 0.0im)
+        @test all(norm.(ρs[:, 1, 1] .- complex_dat[1:ntimes+1, 1, 1]) .< 1e-2)
+        @test all(norm.(ρs[:, 1, 2] .- complex_dat[1:ntimes+1, 2, 1]) .< 1e-2)
+    end
+
     @testset "TEMPO" begin
         t, ρs = TEMPO.propagate(; fbU, Jw=[Jw], β, ρ0, dt, ntimes, kmax, svec, extraargs=TEMPO.TEMPOArgs(; cutoff=1e-15, maxdim=1000))