QMC code restructuring

lib/BloqadeQMC/src/datastructures/Rydberg.jl

@@ -0,0 +1,158 @@
+using Base.Iterators
+
+export Rydberg
+
+# changed type tree s.t. AbstractRydberg is a direct subtype of Hamiltonian... ACTUALLY NO NEED FOR ABSTRACT AT ALL?
+# abstract type AbstractRydberg{O <: AbstractOperatorSampler} <: Hamiltonian{2,O} end
+
+struct Rydberg{O,M <: UpperTriangular{Float64},UΩ <: AbstractVector{Float64}, Uδ <: AbstractVector{Float64}, L <: Lattice} # <: AbstractRydberg{O}
+    op_sampler::O
+    V::M
+    Ω::UΩ
+    δ::Uδ
+    lattice::L
+    energy_shift::Float64
+end
+
+nspins(H::Rydberg) = nspins(H.lattice)
+
+
+function make_prob_vector(H::Type{<:Rydberg}, V::UpperTriangular{T}, Ω::AbstractVector{T}, δ::AbstractVector{T}; epsilon=0.0) where T
+    @assert length(Ω) == length(δ) == size(V, 1) == size(V, 2)
+    @assert (0.0 <= epsilon <= 1.0) "epsilon must be in the range [0, 1]!"
+
+    ops = Vector{NTuple{OP_SIZE, Int}}()
+    p = Vector{T}()
+    energy_shift = zero(T)
+
+    for i in eachindex(Ω)
+        if !iszero(Ω[i])
+            push!(ops, makediagonalsiteop(AbstractLTFIM, i))
+            push!(p, Ω[i] / 2)
+            energy_shift += Ω[i] / 2
+        end
+    end
+
+    Ns = length(Ω)
+    bond_spins = Set{NTuple{2,Int}}()
+    coordination_numbers = zeros(Int, Ns)
+    for j in axes(V, 2), i in axes(V, 1)
+        if i < j && !iszero(V[i, j])
+            push!(bond_spins, (i, j))
+            coordination_numbers[i] += 1
+            coordination_numbers[j] += 1
+        end
+    end
+
+    n = diagonaloperator(H)
+    I = Diagonal(LinearAlgebra.I, 2)
+
+    # TODO: add fictitious bonds if there's a z-field on an "unbonded" site
+    for (site1, site2) in bond_spins
+        # by this point we can assume site1 < site2
+        δb1 = δ[site1] / coordination_numbers[site1]
+        δb2 = δ[site2] / coordination_numbers[site2]
+        local_H = V[site1, site2]*kron(n, n) - δb1*kron(n, I) - δb2*kron(I, n)
+
+        p_spins = -diag(local_H)
+        C = abs(min(0, minimum(p_spins))) + epsilon*abs(minimum(p_spins[2:end]))
+        #dont use the zero matrix element for the epsilon shift
+        p_spins .+= C
+        energy_shift += C
+
+        for (t, p_t) in enumerate(p_spins)
+            push!(p, p_t)
+            push!(ops, (2, t, length(p), site1, site2))
+        end
+    end
+
+    return ops, p, energy_shift
+end
+
+
+###############################################################################
+
+# function Rydberg(dims::NTuple{D, Int}, R_b, Ω, δ; pbc=true, trunc::Int=0, epsilon::Float64=0) where D
+#     if D == 1
+#         lat = Chain(dims[1], 1.0, pbc)
+#     elseif D == 2
+#         lat = Rectangle(dims[1], dims[2], 1.0, 1.0, pbc)
+#     else
+#         error("Unsupported number of dimensions. 1- and 2-dimensional lattices are supported.")
+#     end
+#     return Rydberg(lat, R_b, Ω, δ; trunc=trunc, epsilon=epsilon)
+# end
+
+
+# We  only want one Rydberg() interface, integrated into BloqadeExpr. 
+# Probably, we'll want one function generate_interaction_matrix() that creates V and then feeds that matrix into Rydberg(). 
+# Check with Phil how that function is coming along.
+
+function Rydberg(lat::Lattice, R_b::Float64, Ω::Float64, δ::Float64; trunc::Int=0, epsilon=0)
+    println("Hello!")
+    Ns = nspins(lat)
+    V = zeros(Float64, Ns, Ns)
+
+    if trunc > 0
+        _dist = sort!(collect(Set(lat.distance_matrix)))
+        uniq_dist = Vector{Float64}(undef, 0)
+        for i in eachindex(_dist)
+            if length(uniq_dist) == 0
+                push!(uniq_dist, _dist[i])
+            elseif !(last(uniq_dist) ≈ _dist[i])
+                push!(uniq_dist, _dist[i])
+            end
+        end
+        smallest_k = sort!(uniq_dist)[1:(trunc+1)]
+        dist = copy(lat.distance_matrix)
+        for i in eachindex(dist)
+            if dist[i] > last(smallest_k) && !(dist[i] ≈ last(smallest_k))
+                dist[i] = zero(dist[i])
+            end
+        end
+    elseif lat isa Rectangle && all(lat.PBC)
+        V = zeros(Ns, Ns)
+        K = 3  # figure out an efficient way to set this dynamically
+
+        dist = zeros(Ns, Ns)
+        for v2 in -K:K, v1 in -K:K
+            dV = zeros(Ns, Ns)
+            for x2 in axes(dV, 2), x1 in axes(dV, 1)
+                i1, j1 = divrem(x1 - 1, lat.n1)
+                i2, j2 = divrem(x2 - 1, lat.n1)
+                r = [i2 - i1 + v1*lat.n1, j2 - j1 + v2*lat.n2]
+                dV[x1, x2] += Ω * (R_b/norm(r, 2))^6
+            end
+            # @show v2, v1, maximum(abs, dV)
+
+            V += dV
+        end
+
+        V = (V + V') / 2  # should already be symmetric but just in case
+
+        return Rydberg(UpperTriangular(triu!(V, 1)), Ω*ones(Ns), δ*ones(Ns), lat; epsilon=epsilon)
+    else
+        dist = lat.distance_matrix
+    end
+
+    @inbounds for i in 1:(Ns-1)
+        for j in (i+1):Ns
+            # a zero entry in distance_matrix means there should be no bond
+            V[i, j] = dist[i, j] != 0.0 ? Ω * (R_b / dist[i, j])^6 : 0.0
+        end
+    end
+    V = UpperTriangular(triu!(V, 1))
+
+    return Rydberg(V, Ω*ones(Ns), δ*ones(Ns), lat; epsilon=epsilon)
+end
+
+function Rydberg(V::AbstractMatrix{T}, Ω::AbstractVector{T}, δ::AbstractVector{T}, lattice::Lattice; epsilon=zero(T)) where T
+    ops, p, energy_shift = make_prob_vector(Rydberg, V, Ω, δ, epsilon=epsilon)
+    op_sampler = ImprovedOperatorSampler(AbstractLTFIM, ops, p)
+    return Rydberg{typeof(op_sampler), typeof(V), typeof(Ω), typeof(δ), typeof(lattice)}(op_sampler, V, Ω, δ, lattice, energy_shift)
+end
+
+# Check whether the two functions below are needed in updates.
+
+total_hx(H::Rydberg)::Float64 = sum(H.Ω) / 2
+haslongitudinalfield(H::Rydberg) = !iszero(H.δ)

lib/BloqadeQMC/src/datastructures/alias.jl

@@ -0,0 +1,98 @@
+###############################################################################
+# Walker's Alias Method: draw samples in O(1) time, initialize vector in O(n)
+#   caveat: unsure if it's possible to build an efficient update scheme
+#           for these types of probability vectors
+
+probability_vector(p::Vector{T}) where T = ProbabilityAlias(p)
+
+struct ProbabilityAlias{T} # <: AbstractProbabilityVector{T}     We should be fine without abstract subtyping   
+    normalization::T
+
+    cutoffs::Vector{T}
+    alias::Vector{Int}
+
+    # initialize using Vose's algorithm
+    function ProbabilityAlias{T}(weights::Vector{T}) where {T <: AbstractFloat}
+        if length(weights) == 0
+            throw(ArgumentError("weight vector must have non-zero length!"))
+        end
+        if any(x -> x < zero(T), weights)
+            throw(ArgumentError("weights must be non-negative!"))
+        end
+
+        weights_ = copy(weights)
+        normalization = sum(weights)
+        N = length(weights)
+        avg = normalization / N
+
+        underfull = Stack{Int}()
+        overfull = Stack{Int}()
+
+        for (i, w) in enumerate(weights_)
+            if w >= avg
+                push!(overfull, i)
+            else
+                push!(underfull, i)
+            end
+        end
+
+        cutoffs = zeros(float(T), N)
+        alias = zeros(Int, N)
+
+        while !isempty(underfull) && !isempty(overfull)
+            less = pop!(underfull)
+            more = pop!(overfull)
+
+            cutoffs[less] = weights_[less] * N
+            alias[less] = more
+
+            weights_[more] += weights_[less]
+            weights_[more] -= avg
+
+            if weights_[more] >= avg
+                push!(overfull, more)
+            else
+                push!(underfull, more)
+            end
+        end
+
+        while !isempty(underfull)
+            cutoffs[pop!(underfull)] = normalization
+        end
+
+        while !isempty(overfull)
+            cutoffs[pop!(overfull)] = normalization
+        end
+
+        cutoffs /= normalization
+
+        new{T}(sum(weights), cutoffs, alias)
+    end
+end
+
+ProbabilityAlias(p::Vector{T}) where T = ProbabilityAlias{T}(p)
+@inline length(pvec::ProbabilityAlias) = length(pvec.cutoffs)
+@inline normalization(pvec::ProbabilityAlias) = pvec.normalization
+
+# For safety, copying the following over from probabilityvector.jl. Probably not needed:
+
+# firstindex(::AbstractProbabilityVector) = 1
+# lastindex(pvec::AbstractProbabilityVector) = length(pvec)
+# size(pvec::AbstractProbabilityVector) = (length(pvec),)
+
+function show(io::IO, p::ProbabilityAlias{T}) where T
+    r = repr(p.normalization; context=IOContext(io, :limit=>true))
+    print(io, "ProbabilityAlias{$T}($r)")
+end
+
+# function Base.rand(rng::AbstractRNG, pvec::ProbabilityAlias{T}) where T
+#     u, i::Int = modf(muladd(length(pvec), rand(rng), 1.0))
+#     return @inbounds (u < pvec.cutoffs[i]) ? i : pvec.alias[i]
+# end
+
+# xorshift prngs seem to be noticably faster if you just sample from them twice
+function Base.rand(rng::AbstractRNG, pvec::ProbabilityAlias{T}) where T
+    u = rand(rng)
+    i = rand(rng, 1:length(pvec))
+    return @inbounds (u < pvec.cutoffs[i]) ? i : pvec.alias[i]
+end

lib/BloqadeQMC/src/datastructures/hamiltonian.jl

@@ -0,0 +1,52 @@
+# To befirst merged with Rydberg (the only specific Hamiltonian we need) and eventually integrated with BloqadeExpr interfaces
+abstract type Hamiltonian{D,O<:AbstractOperatorSampler} end
+
+localdim(::Hamiltonian{D}) where {D} = D
+
+zero(H::Hamiltonian{2}) = zeros(Bool, nspins(H))
+zero(H::Hamiltonian) = zeros(Int, nspins(H))
+one(H::Hamiltonian{2}) = ones(Bool, nspins(H))
+one(H::Hamiltonian) = ones(Int, nspins(H))
+
+# nspins(H::Hamiltonian) = H.Ns                         # Rydberg has its own version
+nbonds(H::Hamiltonian) = H.Nb                           # Is this needed for Rydberg?
+
+@inline isdiagonal(H) = op -> isdiagonal(H, op)
+@inline isidentity(H) = op -> isidentity(H, op)
+@inline issiteoperator(H) = op -> issiteoperator(H, op)
+@inline isbondoperator(H) = op -> isbondoperator(H, op)
+@inline getsite(H) = op -> getsite(H, op)
+@inline getbondsites(H) = op -> getbondsites(H, op)
+
+@inline diag_update_normalization(H::Hamiltonian) = normalization(H.op_sampler)
+
+###############################################################################
+
+# These energy functions are only used in Rydberg ED not in LTFIM QMC 
+
+energy(::Type{<:BinaryThermalState}, H::Hamiltonian, β::Float64, n::Int) = H.energy_shift - (n / β)
+function energy(::Type{<:BinaryThermalState}, H::Hamiltonian, β::Float64, ns::Vector{T}) where {T <: Real}
+    E = -mean_and_stderr(ns) / β
+    return H.energy_shift + E
+end
+energy(qmc_state::BinaryQMCState, args...; kwargs...) = energy(typeof(qmc_state), args...; kwargs...)
+
+energy_density(S::Type{<:BinaryQMCState}, H::Hamiltonian, args...; kwargs...) =
+    energy(S, H, args...; kwargs...) / nspins(H)
+
+energy_density(qmc_state::BinaryQMCState, args...; kwargs...) = energy_density(typeof(qmc_state), args...; kwargs...)
+
+###############################################################################
+
+# Move this stuff to state files?
+
+function BinaryGroundState(H::Hamiltonian{2,O}, M::Int, trialstate::Union{Nothing, AbstractTrialState}=nothing) where {K, O <: AbstractOperatorSampler{K}}
+    z = zero(H)         # Here we are initializing state into all zeros
+    BinaryGroundState(z, init_op_list(2*M, Val{K}()), trialstate)::BinaryGroundState{K, typeof(z)}
+end
+
+
+function BinaryThermalState(H::Hamiltonian{2,O}, cutoff::Int) where {K, O <: AbstractOperatorSampler{K}}    #    cutoff is essentially 2M but here we tune it via an algorithm
+    z = zero(H)
+    BinaryThermalState(z, init_op_list(cutoff, Val{K}()))::BinaryThermalState{K, typeof(z)}
+end

lib/BloqadeQMC/src/datastructures/improved_op_sampler.jl

@@ -0,0 +1,73 @@
+# Todo: Get rid of "improved" prefix
+
+# abstract type AbstractOperatorSampler{K, T, P <: AbstractProbabilityVector{T}} end            We probably don't need Abstract type
+firstindex(::ImprovedOperatorSampler) = 1               # changed argument from taking abstract type to concrete type
+lastindex(os::ImprovedOperatorSampler) = length(os)
+@inline normalization(os::AbstractOperatorSampler) = normalization(os.pvec)
+
+# abstract type AbstractImprovedOperatorSampler{K, T, P} <: AbstractOperatorSampler{K, T, P} end        Again, probably no need for Abstract type
+
+struct ImprovedOperatorSampler{K, T, P} # <: AbstractImprovedOperatorSampler{K, T, P}                   no abstract subtyping
+    operators::Vector{NTuple{K, Int}}
+    pvec::ProbabilityAlias              # replaced P with ProbabilityAlias 
+    op_log_weights::Vector{T}
+end
+
+# What do I do with AbstractOperatorSampler in argument here?
+function ImprovedOperatorSampler(H::Type{<:Hamiltonian{2, <:AbstractOperatorSampler}}, operators::Vector{NTuple{K, Int}}, p::Vector{T}) where {T <: AbstractFloat, K}
+    @assert length(operators) == length(p) "Given vectors must have the same length!"
+
+    op_log_weights = log.(p)
+
+    max_mel_ops = Vector{NTuple{K, Int}}()
+    p_modified = Vector{T}()
+
+    # fill with all the site operators first
+    for (i, op) in enumerate(operators)
+        if issiteoperator(H, op) && isdiagonal(H, op)
+            push!(max_mel_ops, op)
+            push!(p_modified, @inbounds p[i])
+        end
+    end
+    idx = findall(isbondoperator(H), operators)
+    ops = operators[idx]
+    p_mod = p[idx]
+
+    perm = sortperm(ops, by=getbondsites(H))
+    ops = ops[perm]
+    p_mod = p_mod[perm]
+
+    op_groups = Dict{NTuple{2, Int}, Vector{NTuple{K, Int}}}()
+    p_groups = Dict{NTuple{2, Int}, Vector{T}}()
+
+    while !isempty(ops)
+        op = pop!(ops)
+        site1, site2 = getbondsites(H, op)
+        p_ = pop!(p_mod)
+
+        if haskey(op_groups, (site1, site2))
+            push!(op_groups[(site1, site2)], op)
+            push!(p_groups[(site1, site2)], p_)
+        else
+            op_groups[(site1, site2)] = [op]
+            p_groups[(site1, site2)] = [p_]
+        end
+    end
+
+    for (k, p_gr) in p_groups
+        i = argmax(p_gr)
+        push!(max_mel_ops, op_groups[k][i])
+        push!(p_modified, p_gr[i])
+    end
+
+    pvec = probability_vector(p_modified)
+    return ImprovedOperatorSampler{K, T, typeof(pvec)}(max_mel_ops, pvec, op_log_weights)                   # Here, we return only the max matrix elements and associated compressed probability vector. Full information still contained in op_log_weights
+end
+
+@inline rand(rng::AbstractRNG, os::ImprovedOperatorSampler) = @inbounds os.operators[rand(rng, os.pvec)]
+
+Base.@propagate_inbounds getlogweight(os::ImprovedOperatorSampler, w::Int) = os.op_log_weights[w]
+
+@inline length(os::ImprovedOperatorSampler) = length(os.operators)
+
+##############################################################################