Take real, imag of Pow function. (#12)

* new subs

* fix taking real for Pow function

example/symad2.jl

@@ -0,0 +1,16 @@
+using Yao, YaoExtensions, Test
+
+@vars α β γ
+c = chain(put(3, 2=>Rx(α)), control(3, 2, 1=>Ry(β)), put(3, (1,2)=>rot(kron(X, X), γ)))
+h = heisenberg(3);
+es = expect(h, zero_state(Basic, 3)=>c)
+gs = expect'(h, zero_state(Basic, 3)=>c).second
+assign = Dict(α=>0.5, β=>0.7, γ=>0.8)
+
+cn = subs(Float64, c, assign...)
+en = expect(h, zero_state(3)=>cn)
+gn = expect'(h, zero_state(3)=>cn).second
+es_ = subs(es, assign...)
+gs_ = map(x->subs(x, assign...), gs)
+@test ComplexF64.(gs_) ≈ gn
+@test ComplexF64(es_) ≈ en

src/YaoSym.jl

@@ -1,7 +1,7 @@
 module YaoSym
 
 using SymEngine: @vars, Basic
-export @vars, Basic
+export @vars, Basic, subs
 
 include("register.jl")
 include("symengine/backend.jl")

src/symengine/blocks.jl

@@ -7,6 +7,9 @@ const BiVarOp = Union{[SymEngine.BasicType{Val{i}} for i in op_types]...}
 export smat, @vars
 smat(block::AbstractBlock) = mat(Basic, block)
 
+simag = SymFunction("Im")
+sreal = SymFunction("Re")
+
 Base.promote_rule(::Type{Bool}, ::Type{Basic}) = Basic
 Base.conj(x::Basic) = Basic(conj(SymEngine.BasicType(x)))
 Base.conj(x::BasicType) = real(x) - im*imag(x)
@@ -32,6 +35,24 @@ function Base.imag(x::BasicType{Val{:Mul}})
     #return res
 end
 
+function Base.real(x::BasicType{Val{:Pow}})
+    a, b = get_args(x.x)
+    if imag(a) == 0 && imag(b) == 0
+        return x.x
+    else
+        return sreal(x.x)
+    end
+end
+
+function Base.imag(x::BasicType{Val{:Pow}})
+    a, b = get_args(x.x)
+    if imag(a) == 0 && imag(b) == 0
+        return Basic(0)
+    else
+        return simag(x.x)
+    end
+end
+
 function Base.real(x::BasicType{Val{:Mul}})
     args = (get_args(x.x)...,)
     get_mul_real(args)
@@ -44,7 +65,7 @@ function get_mul_imag(args::NTuple{N,Any}) where N
 end
 get_mul_imag(args::Tuple{Basic}) = imag(args[1])
 
-function get_mul_real(args::Tuple{N,Any}) where N
+function get_mul_real(args::NTuple{N,Any}) where N
     real(args[1])*get_mul_real(args[2:end]) - imag(args[1])*get_mul_imag(args[2:end])
 end
 get_mul_real(args::Tuple{Basic}) = real(args[1])
@@ -54,7 +75,7 @@ function Base.real(x::BasicTrigFunction)
     if imag(a) == 0
         return x.x
     else
-        error("The imag of triangular function $(typeof(x)) with complex argument is not yet implemented!")
+        return sreal(x.x)
     end
 end
 
@@ -63,7 +84,7 @@ function Base.imag(x::BasicTrigFunction)
     if imag(a) == 0
         return Basic(0)
     else
-        error("The imag of triangular function $(typeof(x)) with complex argument is not yet implemented!")
+        return simag(x.x)
     end
 end
 
@@ -101,3 +122,16 @@ YaoBlocks.PSwap{N}(locs::Tuple{Int, Int}, θ::SymReal) where N = YaoBlocks.PutBl
 
 YaoBlocks.pswap(n::Int, i::Int, j::Int, α::SymReal) = PSwap{n}((i,j), α)
 YaoBlocks.pswap(i::Int, j::Int, α::SymReal) = n->pswap(n,i,j,α)
+
+export subs, chiparams
+chiparams(blk::RotationGate, param) = rot(blk.block, param)
+chiparams(blk::ShiftGate, param) = shift(param)
+chiparams(blk::PhaseGate, param) = phase(param)
+chiparams(blk::TimeEvolution, param) = time_evolve(blk.H, param, tol=blk.tol)
+chiparams(blk::AbstractBlock, params...) = niparams(blk) == length(params) == 0 ? blk : throw(NotImplementedError(:chiparams, (blk, params...)))
+
+SymEngine.subs(c::AbstractBlock, args...; kwargs...) = subs(Basic, c, args...; kwargs...)
+function SymEngine.subs(::Type{T}, c::AbstractBlock, args...; kwargs...) where T
+    c = chiparams(c, map(x->T(subs(x, args...; kwargs...)), getiparams(c))...)
+    chsubblocks(c, [subs(T, blk, args..., kwargs...) for blk in subblocks(c)])
+end

src/symengine/register.jl

@@ -1,6 +1,7 @@
 using YaoBase, SparseArrays, BitBasis, YaoArrayRegister, SymEngine
 export @ket_str, @bra_str
 export SymReg, AdjointSymReg, SymRegOrAdjointSymReg, expand
+export szero_state
 
 YaoArrayRegister._warn_type(raw::AbstractArray{Basic}) = nothing
 
@@ -34,6 +35,7 @@ Base.:(*)(x::AdjointSymReg{B, MT}, y::AdjointSymReg{B, MT}) where {B, MT} = adjo
 Base.:(^)(x::AdjointSymReg{B, MT}, n::Int) where {B, MT} = adjoint(parent(x)^n)
 
 SymEngine.expand(x::SymReg{B}) where B = ArrayReg{B}(expand.(state(x)))
+szero_state(args...; kwargs...) = zero_state(Basic, args...; kwargs...)
 
 function YaoBase.partial_tr(r::SymReg{B}, locs) where B
     orders = setdiff(1:nqubits(r), locs)

test/symengine/blocks.jl

@@ -4,11 +4,13 @@ using SymEngine
 using Test
 
 @testset "imag and conj" begin
-    @vars a b c
+    @vars a b c x
     @test imag(Basic(2+im*a)) == a
     @test imag(Basic((2+im)*(im*a+1))) == 2*a + 1
     @test real(sin(a)) == sin(a)
     @test imag(sin(a)) == 0
+    @test imag(x^2) == 0
+    @test real(x^2) == x^2
 end
 
 @testset "mat" begin
@@ -41,6 +43,17 @@ end
     @test pswap(4,2,1,θ) isa PSwap{4,Basic}
 end
 
+@testset "sub" begin
+    @vars θ ϕ
+    @test subs(Float64, Rx(θ), θ=>0.5) == Rx(0.5)
+    @test subs(Float64, shift(θ), θ=>0.5) == shift(0.5)
+    @test subs(Float64, phase(θ), θ=>0.5) == phase(0.5)
+    @test subs(Float64, chain(control(2,1,2=>shift(θ)), put(2,1=>chain(Rx(θ), kron(Rx(θ))))), θ=>0.5) ==
+        chain(control(2,1,2=>shift(0.5)), put(2,1=>chain(Rx(0.5), kron(Rx(0.5)))))
+    @test subs(Float64, time_evolve(X, θ), θ=>0.5) == time_evolve(X, 0.5)
+    @test subs(Rx(θ), θ=>ϕ) == Rx(ϕ)
+end
+
 @testset "Sym Engine Conj" begin
     @vars a b
     @test cos(exp(a+im*b))' == cos(exp(a-im*b))