make terms primtive block

lib/EaRydCUDA/Project.toml

@@ -21,7 +21,7 @@ EaRydODE = "0.1"
 OrdinaryDiffEq = "5"
 Reexport = "1"
 EaRydCore = "0.1"
-Yao = "0.6"
+Yao = "0.7"
 julia = "1.6"
 
 [extras]

lib/EaRydCore/Project.toml

@@ -46,7 +46,7 @@ StatsBase = "0.33"
 ThreadsX = "0.1"
 Transducers = "0.4"
 Unitful = "1"
-Yao = "0.6"
+Yao = "0.7"
 julia = "1.4"
 
 [extras]

lib/EaRydCore/src/hamiltonian/cache.jl

@@ -0,0 +1,130 @@
+function get_matrix(::Type{Tv}, op::AbstractBlock, ::FullSpace) where Tv
+    return mat(Tv, op)
+end
+
+function get_matrix(::Type{Tv}, op::AbstractBlock, space::Subspace) where Tv
+    return mat(Tv, op, space)
+end
+
+function get_matrix(::Type{Tv}, op::AbstractTerm, space::AbstractSpace) where Tv
+    return SparseMatrixCSC{Tv}(op, space)
+end
+
+struct ConstTermCache{FS <: Tuple, HS <: Tuple}
+    fs::FS # time-dependent factors
+    hs::HS # const terms
+end
+
+function storage_size(h::ConstTermCache)
+    return sum(storage_size, h.hs)
+end
+
+# split const term and its dynamic prefactors from hamiltonian expr
+function split_const_term(::Type{Tv}, h::Hamiltonian, space::AbstractSpace) where {Tv}
+    fs, hs = [], []
+    for t in h.terms, (f, h) in _split_term(Tv, t, space)
+        push!(fs, f)
+        # NOTE: we force converting blocks to a matrix as a workaround
+        # of https://github.com/QuantumBFS/BQCESubroutine.jl/issues/37
+        # so that we don't need to special case blocks to preallocate
+        # the intermediate state for dstate.
+        if h isa AbstractBlock
+            push!(hs, get_matrix(Tv, h, space))
+        else
+            push!(hs, h)
+        end
+    end
+    return ConstTermCache((fs...,), (hs...,))
+end
+
+function _split_term(::Type{Tv}, h::RydInteract, space::AbstractSpace) where {Tv}
+    # TODO: actually implement it as Diagonal
+    ((_const_param_, Diagonal(Vector(diag(SparseMatrixCSC{Tv}(h, space))))), )
+end
+
+function _split_term(::Type{Tv}, h::Negative, space::AbstractSpace) where {Tv}
+    return map(_split_term(Tv, h.term, space)) do (f, h)
+        f, -h
+    end
+end
+
+_const_param_(t) = one(t)
+
+function _split_term(::Type{Tv}, h::XTerm, space::AbstractSpace) where {Tv}
+    n = nsites(h)
+    @switch (h.Ωs, h.ϕs) begin
+        @case (Ωs::ConstParamListType, ϕ::Number) || (Ωs::ConstParamListType, ::Nothing) || (Ω::Number, ϕ::Number) ||
+        (Ω::Number, ::ConstParamListType) || (Ω::Number, ::Nothing)
+            ((_const_param_, SparseMatrixCSC{Tv, Cint}(h, space)), )
+        @case (Ωs::AbstractVector, ϕs::ConstParamListType) # directly apply is faster
+            map(enumerate(zip(Ωs, ϕs))) do (i, (Ω, ϕ))
+                x_phase = PermMatrix([2, 1], Tv[exp(ϕ * im), exp(-ϕ * im)])
+                t->Ω(t)/2, put(n, i => matblock(x_phase))
+            end
+        @case (Ωs::ConstParamListType, ϕs::ParamsList) # directly apply is faster
+            op1 = map(enumerate(zip(Ωs, ϕs))) do (i, (Ω, ϕ))
+                t->(Ω/2 * exp(ϕ(t) * im)), put(n, i => matblock(Tv[0 1;0 0]))
+            end
+
+            op2 = map(enumerate(zip(Ωs, ϕs))) do (i, (Ω, ϕ))
+                t->(Ω/2 * exp(-ϕ(t) * im)), put(n, i => matblock(Tv[0 0;1 0]))
+            end
+            return (op1..., op2...)
+        @case (Ωs::ParamsList, ϕs::ParamsList)
+            op1 = map(enumerate(zip(Ωs, ϕs))) do (i, (Ω, ϕ))
+                t->(Ω(t)/2 * exp(ϕ(t) * im)), put(n, i => matblock(Tv[0 1;0 0]))
+            end
+
+            op2 = map(enumerate(zip(Ωs, ϕs))) do (i, (Ω, ϕ))
+                t->(Ω(t)/2 * exp(-ϕ(t) * im)), put(n, i => matblock(Tv[0 0;1 0]))
+            end
+            return (op1..., op2...)
+        @case (Ωs::ConstParamListType, ϕ)
+            op1 = map(enumerate(zip(Ωs, ϕs))) do (i, (Ω, ϕ))
+                t->(Ω/2 * exp(ϕ(t) * im)), put(n, i => matblock(Tv[0 1;0 0]))
+            end
+
+            op2 = map(enumerate(zip(Ωs, ϕs))) do (i, (Ω, ϕ))
+                t->(Ω/2 * exp(-ϕ(t) * im)), put(n, i => matblock(Tv[0 0;1 0]))
+            end
+            return (op1..., op2...)
+        @case (Ωs::ParamsList, ::Nothing)
+            map(enumerate(Ωs)) do (i, Ω)
+                t->Ω(t)/2, put(n, i=>X)
+            end
+        @case (Ω::Number, ::ParamsList)
+            op1 = map(enumerate(ϕs)) do (i, ϕ)
+                t->(Ω/2 * exp(ϕ(t) * im)), put(n, i => matblock(Tv[0 1;0 0]))
+            end
+
+            op2 = map(enumerate(ϕs)) do (i, ϕ)
+                t->(Ω/2 * exp(-ϕ(t) * im)), put(n, i => matblock(Tv[0 0;1 0]))
+            end
+            return (op1..., op2...)
+        @case (Ω, ϕ::Number)
+            A = get_matrix(Tv, sum(put(n, i=>matblock(Tv[0 1;0 0]))), space)
+            B = get_matrix(Tv, sum(put(n, i=>matblock(Tv[0 0;1 0]))), space)
+            return (t->Ω(t)/2 * exp(ϕ * im), A), (t->Ω(t)/2 * exp(-ϕ * im), B)
+        @case (Ω, ::Nothing) # no 1/2 in prefactor, it's in the matrix already
+            return ((t->Ω(t), SparseMatrixCSC{Tv, Cint}(XTerm(n, 1.0), space)), )
+        @case (Ω, ϕ)
+            A = get_matrix(Tv, sum(put(n, i=>matblock(Tv[0 1;0 0]))), space)
+            B = get_matrix(Tv, sum(put(n, i=>matblock(Tv[0 0;1 0]))), space)
+            return (t->Ω(t)/2 * exp(ϕ(t) * im), A), (t->Ω(t)/2 * exp(-ϕ(t) * im), B)
+    end
+end
+
+function _split_term(::Type{Tv}, h::NTerm, space::AbstractSpace) where {Tv}
+    n = nsites(h)
+    return if h.Δs isa ConstParamType
+        M = Diagonal(Vector(diag(SparseMatrixCSC{Tv}(h, space))))
+        ((_const_param_, M), )
+    elseif h.Δs isa ParamsList
+        return map(enumerate(h.Δs)) do (i, Δ)
+            Δ, put(n, i=>Yao.ConstGate.P1)
+        end
+    else
+        M = Diagonal(Vector(diag(SparseMatrixCSC{Tv}(NTerm(n, one(Tv)), space))))
+        return ((h.Δs, M), )
+    end
+end

lib/EaRydCore/src/hamiltonian/hamiltonian.jl

@@ -5,3 +5,4 @@ include("operations.jl")
 include("sparse.jl")
 include("interface.jl")
 include("adapt.jl")
+include("cache.jl")

lib/EaRydCore/src/hamiltonian/operations.jl

@@ -1,31 +1,7 @@
-Base.:(+)(x::AbstractTerm, y::AbstractTerm) = Hamiltonian((x, y))
-Base.:(+)(x::AbstractTerm, y::Hamiltonian) = Hamiltonian((x, y.terms...))
-Base.:(+)(x::Hamiltonian, y::AbstractTerm) = Hamiltonian((x.terms..., y))
-Base.:(+)(x::Hamiltonian, y::Hamiltonian) = Hamiltonian((x.terms..., y.terms...))
-
-# absorb - to RHS
-Base.:(-)(x::AbstractTerm, y::AbstractTerm) = Hamiltonian((x, -y))
-Base.:(-)(x::AbstractTerm, y::Hamiltonian) = Hamiltonian((x, map(-, y.terms)...))
-Base.:(-)(x::Hamiltonian, y::AbstractTerm) = Hamiltonian((x.terms..., -y))
-Base.:(-)(x::Hamiltonian, y::Hamiltonian) = Hamiltonian((x.terms..., map(-, y.terms)...))
-
-# Base.:(-)(x::XTerm{<:ConstParamType}) = XTerm(x.nsites, map(-, x.Ωs), x.ϕs)
-# Base.:(-)(x::XTerm) = XTerm(x.nsites, t->-x.Ωs(t), x.ϕs)
-
-# Base.:(-)(x::ZTerm{<:ConstParamType}) = ZTerm(x.nsites, map(-, x.Δs))
-# Base.:(-)(x::ZTerm) = ZTerm(x.nsites, t->-x.Δs(t))
-# Base.:(-)(x::NTerm{<:ConstParamType}) = NTerm(x.nsites, map(-, x.Δs))
-# Base.:(-)(x::NTerm) = NTerm(x.nsites, t->-x.Δs(t))
-
-
 function Base.:(==)(x::RydInteract, y::RydInteract)
     return (x.atoms == y.atoms) && (x.C == y.C)
 end
 
-function Base.:(==)(x::Hamiltonian, y::Hamiltonian)
-    return all(t in y.terms for t in x.terms) && all(t in x.terms for t in y.terms)
-end
-
 """
     getterm(terms, k, k_site)
 
@@ -64,9 +40,5 @@ function getterm(t::NTerm, k, k_site)
     end
 end
 
-function getterm(t::Hamiltonian, k, k_site)
-    error("composite Hamiltonian term cannot be indexed")
-end
-
 space_size(term::AbstractTerm, s::FullSpace) = 1 << nsites(term)
 space_size(::AbstractTerm, s::Subspace) = length(s)

lib/EaRydCore/src/hamiltonian/printing.jl

@@ -55,19 +55,6 @@ function print_term(io::IO, t::RydInteract)
     printstyled(io, "n_i n_j", color=:light_blue)
 end
 
-function print_term(io::IO, t::Hamiltonian)
-    indent = get(io, :indent, 2)
-    for (i, each) in enumerate(t.terms)
-        println(io, " "^indent, " Term ", i)
-        print_term(IOContext(io, :indent=>indent+2), each)
-
-        if i != lastindex(t.terms)
-            println(io)
-            println(io)
-        end
-    end
-end
-
 function print_term(io::IO, t::Negative)
     print_term(IOContext(io, :negative=>true), t.term)
 end

lib/EaRydCore/src/hamiltonian/sparse.jl

@@ -6,9 +6,9 @@ function SparseArrays.SparseMatrixCSC{Tv, Ti}(term::AbstractTerm, s::AbstractSpa
     return H
 end
 
-function SparseArrays.SparseMatrixCSC{Tv, Ti}(term::Hamiltonian, s::AbstractSpace=fullspace) where {Tv, Ti}
-    return sum(SparseMatrixCSC{Tv, Ti}(t, s) for t in term.terms)
-end
+# function SparseArrays.SparseMatrixCSC{Tv, Ti}(term::Hamiltonian, s::AbstractSpace=fullspace) where {Tv, Ti}
+#     return sum(SparseMatrixCSC{Tv, Ti}(t, s) for t in term.terms)
+# end
 
 function SparseArrays.SparseMatrixCSC{Tv, Ti}(t::Negative, s::AbstractSpace=fullspace) where {Tv, Ti}
     H = SparseMatrixCSC{Tv, Ti}(t.term, s)

lib/EaRydCore/src/hamiltonian/types.jl

@@ -3,7 +3,7 @@
 
 Abstract term for hamiltonian terms.
 """
-abstract type AbstractTerm end
+abstract type AbstractTerm <: PrimitiveBlock{2} end
 
 const ConstParamType = Union{Number, AbstractVector{<:Number}, NTuple{N, <:Number} where N}
 const ConstParamListType = Union{AbstractVector{<:Number}, NTuple{N, <:Number} where N}
@@ -143,31 +143,6 @@ function NTerm(nsites::Int, Δs)
     NTerm{typeof(Δs)}(nsites, Δs)
 end
 
-struct Negative{Term} <: AbstractTerm
-    term::Term
-end
-
-const OrNegative{T} = Union{Negative{T}, T}
-
-Base.:(-)(t::AbstractTerm) = Negative(t)
-Base.:(-)(t::Negative) = t.term
-
-struct Hamiltonian{Terms <: Tuple} <: AbstractTerm
-    terms::Terms
-
-    function Hamiltonian(terms::Terms) where {Terms}
-        first_nsites = nsites(first(terms))
-        for idx in 2:length(terms)
-            first_nsites == nsites(terms[idx]) ||
-                throw(ArgumentError(
-                    "nsites mismatch, " *
-                    "expect $first_nsites, got $(nsites(terms[idx]))"
-                ))
-        end
-        new{Terms}(terms)
-    end
-end
-
 # try to infer number of sites from the input
 # this is only necessary for CUDA
 to_tuple(xs) = (xs..., ) # make it type stable
@@ -239,10 +214,10 @@ function nsites end
 nsites(t::XTerm) = t.nsites
 nsites(t::ZTerm) = t.nsites
 nsites(t::NTerm) = t.nsites
-nsites(t::Hamiltonian) = nsites(t.terms[1])
-nsites(t::Negative) = nsites(t.term)
 nsites(t::RydInteract) = length(t.atoms)
 
+Yao.nqudits(t::AbstractTerm) = nsites(t)
+
 function nsites(terms::Vector{<:AbstractTerm})
     term_nsites = nsites(first(terms))
     for i in 2:length(terms)
@@ -259,7 +234,6 @@ Base.eltype(t::ZTerm) = eltype(t.Δs)
 Base.eltype(t::NTerm) = eltype(t.Δs)
 Base.eltype(t::Negative) = eltype(t.term)
 Base.eltype(t::RydInteract) = typeof(t.C)
-Base.eltype(t::Hamiltonian) = promote_type(eltype.(t.terms)...)
 
 function Base.isreal(t::XTerm)
     isnothing(t.ϕs) ? true :
@@ -270,12 +244,8 @@ Base.isreal(t::ZTerm) = true
 Base.isreal(t::NTerm) = true
 Base.isreal(t::RydInteract) = true
 Base.isreal(t::Negative) = isreal(t.term)
-Base.isreal(t::Hamiltonian) = all(isreal, t.terms)
 
 Base.iszero(t::XTerm) = iszero(t.Ωs)
 Base.iszero(t::Union{NTerm, ZTerm}) = iszero(t.Δs)
 Base.iszero(t::RydInteract) = iszero(t.C)
 Base.iszero(t::Negative) = iszero(t.term)
-Base.iszero(t::Hamiltonian) = all(iszero, t.terms)
-
-Base.getindex(t::Hamiltonian, i::Int) = t.terms[i]

lib/EaRydCore/src/mat.jl

@@ -12,8 +12,8 @@ YaoBlocks.mat(x::AbstractBlock, s::Subspace) =
 YaoBlocks.mat(::Type{T}, d::TrivialGate{N}, s::Subspace) where {T,N} = IMatrix{length(s),T}()
 YaoBlocks.mat(::Type{T}, pb::PutBlock{N,C}, s::Subspace) where {T,N,C} =
     _cunmat(s.subspace_v, s.map, (), (), mat(T, pb.content), pb.locs)
-YaoBlocks.mat(::Type{T}, c::Subroutine{N,<:AbstractBlock}, s::Subspace) where {N,T} =
-    mat(T, PutBlock{N}(c.content, c.locs), s)
+YaoBlocks.mat(::Type{T}, c::Subroutine{D, <:AbstractBlock}, s::Subspace) where {D,T} =
+    mat(T, PutBlock(c.n, c.content, c.locs), s)
 YaoBlocks.mat(::Type{T}, x::Scale, s::Subspace) where {T} = YaoBlocks.factor(x) * mat(T, content(x), s)
 YaoBlocks.mat(::Type{T}, blk::Daggered, s::Subspace) where {T} = adjoint(mat(T, content(blk), s))
 YaoBlocks.mat(::Type{T}, c::ControlBlock{N,BT,C}, s::Subspace) where {T,N,BT,C} =

lib/EaRydCore/src/measure.jl

@@ -33,3 +33,8 @@ end
 function Yao.measure(; nshots=1)
     reg -> Yao.measure(reg; nshots=nshots)
 end
+
+# TODO: remove this after https://github.com/QuantumBFS/Yao.jl/issues/338
+function Yao.expect(op::AbstractBlock, reg::RydbergReg)
+    return reg' * apply!(copy(reg), op)
+end

lib/EaRydCore/src/mis.jl

@@ -83,7 +83,7 @@ struct ConfigAmplitude{Reg <: Yao.AbstractRegister}
     range::UnitRange{Int}
 end
 
-ConfigAmplitude(reg::Yao.AbstractRegister{1}) = ConfigAmplitude(reg, 1:size(reg.state, 1))
+ConfigAmplitude(reg::Yao.AbstractRegister{2}) = ConfigAmplitude(reg, 1:size(reg.state, 1))
 
 Base.eltype(it::ConfigAmplitude) = Tuple{Int, Yao.datatype(it.reg)}
 Base.length(it::ConfigAmplitude) = length(it.range)

lib/EaRydCore/src/register.jl

@@ -2,7 +2,7 @@ abstract type MemoryLayout end
 struct RealLayout <: MemoryLayout end
 struct ComplexLayout <: MemoryLayout end
 
-struct RydbergReg{Layout <: MemoryLayout, State, Space <: Subspace} <: Yao.AbstractRegister{1}
+struct RydbergReg{Layout <: MemoryLayout, State, Space <: Subspace} <: Yao.AbstractRegister{2}
     natoms::Int
     layout::Layout
     state::State
@@ -176,6 +176,6 @@ function set_zero_state!(r::Yao.ArrayReg)
     return r
 end
 
-function Base.:*(bra::Yao.YaoArrayRegister.AdjointRegister{1, <:RydbergReg}, ket::RydbergReg)
+function Base.:*(bra::Yao.YaoArrayRegister.AdjointRegister{2, <:RydbergReg}, ket::RydbergReg)
     return dot(statevec(parent(bra)), statevec(ket))
 end

lib/EaRydCore/test/cache.jl

@@ -0,0 +1,30 @@
+using Test
+using EaRydCore
+using SparseArrays
+using LinearAlgebra
+using EaRydCore: split_const_term
+
+atoms = square_lattice(4, 0.8)
+
+@testset "split_const_term $(nameof(typeof(space)))" for space in [FullSpace(), blockade_subspace(atoms)]
+    for h in [
+        rydberg_h(atoms; Δ=0.1, Ω=0.1),
+        rydberg_h(atoms; Δ=0.1, Ω=sin),
+        rydberg_h(atoms; Δ=cos, Ω=sin),
+        rydberg_h(atoms; Δ=cos, Ω=[sin, sin, sin, sin]),
+        rydberg_h(atoms; Δ=[cos, cos, cos, cos], Ω=[sin, sin, sin, sin]),
+    ]
+
+        H = SparseMatrixCSC{ComplexF64}(h(0.1), space)
+        tc = split_const_term(ComplexF64, h, space)
+        M = sum(zip(tc.fs, tc.hs)) do (f, h)
+            if h isa AbstractBlock
+                f(0.1) * mat(h)
+            else
+                f(0.1) * h
+            end
+        end
+
+        @test M ≈ H
+    end
+end

lib/EaRydCore/test/instructs.jl

@@ -20,4 +20,4 @@ end
         M = SparseMatrixCSC(mat(g))
         @test expect(g, r) ≈ r.state' * M[ss, ss] * r.state
     end
-end
+end

lib/EaRydCore/test/runtests.jl

@@ -15,6 +15,7 @@ end
 
 @testset "hamiltonian" begin
     include("hamiltonian.jl")
+    include("cache.jl")
 end
 
 @testset "QAOA emulator" begin