ZXW rewrite rules + Differentiation

Project.toml

@@ -11,6 +11,7 @@ MLStyle = "d8e11817-5142-5d16-987a-aa16d5891078"
 Multigraphs = "7ebac608-6c66-46e6-9856-b5f43e107bac"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 YaoHIR = "6769671a-fce8-4286-b3f7-6099e1b1298a"
 YaoLocations = "66df03fb-d475-48f7-b449-3d9064bf085b"
 

src/ZXCalculus.jl

@@ -16,14 +16,15 @@ export SpiderType, EdgeType, ZXWSpiderType
 export AbstractZXDiagram, ZXDiagram, ZXGraph
 export ZXWDiagram
 export Rule, Match
+export ZXWRule
 
 export spider_type, phase, spiders, rem_spider!, rem_spiders!, scalar
 export parameter
-export push_gate!, pushfirst_gate!, tcount
+export push_gate!, pushfirst_gate!, tcount, insert_wtrig!
 export convert_to_chain, convert_to_zxd
 export rewrite!, simplify!, clifford_simplification, full_reduction, 
     circuit_extraction, phase_teleportation
-export random_circuit
+export substitute_variables!, expval_circ!, diff_diagram!, stack_zxwd, concat!, integrate!
 
 include("parameter.jl")
 include("phase.jl")
@@ -40,6 +41,8 @@ include("to_eincode.jl")
 include("utils.jl")
 
 include("rules.jl")
+include("zxw_rules.jl")
+include("diff.jl")
 include("simplify.jl")
 include("circuit_extraction.jl")
 include("phase_teleportation.jl")

src/diff.jl

@@ -0,0 +1,202 @@
+using ZXCalculus: contains, dagger, concat!, stack
+
+"""
+Take derivative of ZXWDiagram with respect to a parameter
+
+Assuming Spiders have Parameter of type PiUnit which is parameterized purely by θ
+"""
+function diff_diagram!(zxwd::ZXWDiagram{T,P}, θ::Symbol) where {T,P}
+    vs_pos = symbol_vertices(zxwd, θ)
+    vs_neg = symbol_vertices(zxwd, θ; neg = true)
+
+    length(vs_pos) + length(vs_neg) == 0 && return zxwd
+
+    add_global_phase!(zxwd, P(π / 2))
+    w_trig_vs = T[]
+
+    for v in vs_pos
+
+        h_v = @match spider_type(zxwd,v) begin
+            X(_) => add_spider!(zxwd, H, [v])
+            Z(_) => v
+        end
+
+        x_v = add_spider!(zxwd, X(Parameter(Val(:PiUnit), 1.0)), [h_v])
+        w_v = add_spider!(zxwd, D, [x_v])
+
+        frac_v = @match parameter(zxwd, v) begin
+            PiUnit(_,_) => w_v
+            Factor(_,_) => error("Only supports PiUnit differentiation")
+            _ => error("not a valid parameter")
+        end
+        push!(w_trig_vs, frac_v)
+    end
+
+    for v in vs_neg
+
+        h_v = @match spider_type(zxwd,v) begin
+            X(_) => add_spider!(zxwd, H, [v])
+            Z(_) => v
+        end
+
+        x_v = add_spider!(zxwd, X(Parameter(Val(:PiUnit), 1.0)), [h_v])
+        w_v = add_spider!(zxwd, D, [x_v])
+
+        frac_v = @match spider_type(zxwd, v).p begin
+            PiUnit(_, _) => add_spider!(zxwd, Z(Parameter(Val(:PiUnit), 1.0)), [w_v])
+            Factor(_, _) => error("Only supports PiUnit differentiation")
+            _ => error("not a valid parameter")
+        end
+        push!(w_trig_vs, frac_v)
+    end
+
+    head = insert_wtrig!(zxwd, w_trig_vs)
+
+    add_spider!(zxwd, X(Parameter(Val(:PiUnit), 1.0)), [head])
+    # our definition of x_tensor exceeds one power of sqrt(2)
+    add_power!(zxwd, -1)
+
+    return zxwd
+end
+
+"""
+Construct ZXW Diagram for representing the expectation value circuit
+"""
+function expval_circ!(zxwd::ZXWDiagram{T,P}, H::String) where {T,P}
+    # convert U to U H U^\dagger
+    zxwd_dag = dagger(zxwd)
+    for (i, h) in enumerate(H)
+        if h == 'Z'
+            push_gate!(zxwd, Val(:Z), i, 1.0)
+        elseif h == 'X'
+            push_gate!(zxwd, Val(:X), i, 1.0)
+        elseif h == 'Y'
+            push_gate!(zxwd, Val(:Z), i, 1.0)
+            push_gate!(zxwd, Val(:X), i, 1.0)
+            add_global_phase!(zxwd, P(π / 2))
+        elseif h == 'I'
+            continue
+        else
+            error("Invalid Hamiltonian, enter only Z, X, Y")
+        end
+    end
+    concat!(zxwd, zxwd_dag)
+    return zxwd
+end
+
+"""
+
+Finds vertices of Spider that contains the parameter θ or -θ
+"""
+function symbol_vertices(zxwd::ZXWDiagram{T,P}, θ::Symbol; neg::Bool = false) where {T,P}
+    if neg
+        target = Expr(:call, :-, θ)
+    else
+        target = θ
+    end
+    matched = T[]
+    for v in vertices(zxwd.mg)
+        res = @match spider_type(zxwd, v) begin
+            Z(p1) && if contains(p1, target)
+            end => v
+            X(p1) && if contains(p1, target)
+            end => v
+            _ => nothing
+        end
+        res !== nothing && push!(matched, v)
+    end
+    return matched
+end
+
+"""
+Replace symbols in ZXW Diagram with specific values
+"""
+function substitute_variables!(
+    zxwd::ZXWDiagram{T,P},
+    sbd::Dict{Symbol,<:Number},
+) where {T,P}
+    for (θ, val) in sbd
+        for negative in [false, true]
+            matched_pos = symbol_vertices(zxwd, θ; neg = negative)
+            val = negative ? -val : val
+            for idx in matched_pos
+                p = spider_type(zxwd, idx).p
+                @match p begin
+                    PiUnit(pu, _) => set_phase!(zxwd, idx, Parameter(Val(:PiUnit), val))
+                    Factor(pf, _) => set_phase!(zxwd, idx, Parameter(Val(:Factor), val))
+                end
+            end
+        end
+    end
+    return zxwd
+end
+
+"""
+Integrate over the Spiders at locs with respect to the parameter θ.
+
+User need to check that the parameters are indeed in the form of k * θ where k is Int
+"""
+function integrate!(zxwd::ZXWDiagram{T,P}, locs::Vector{T}) where {T,P}
+    length(locs) == 2 && return integrate2!(zxwd, locs[1], locs[2])
+    length(locs) == 4 && return integrate4!(zxwd, locs[1], locs[2], locs[3], locs[4])
+end
+
+function integrate2!(zxwd::ZXWDiagram{T,P}, loc1::T, loc2::T) where {T,P}
+    loc1 = int_prep!(zxwd, loc1)
+    loc2 = int_prep!(zxwd, loc2)
+    add_edge!(zxwd.mg, loc1, loc2)
+    return zxwd
+end
+
+"""
+Integrate two pairs of +/- parameter. Theorem 23 of https://arxiv.org/abs/2201.13250
+"""
+function integrate4!(zxwd::ZXWDiagram{T,P}, loca::T, locb::T, locc::T, locd::T) where {T,P}
+    loca = int_prep!(zxwd, loca)
+    locb = int_prep!(zxwd, locb)
+    locc = int_prep!(zxwd, locc)
+    locd = int_prep!(zxwd, locd)
+
+    # a, b = + , - \theta
+    # c, d = + , - \theta
+    loca = add_spider!(zxwd, Z(Parameter(Val(:PiUnit), 0)), [loca])
+    locb = add_spider!(zxwd, Z(Parameter(Val(:PiUnit), 0)), [locb])
+    locc = add_spider!(zxwd, X(Parameter(Val(:PiUnit), 0)), [locc])
+    locd = add_spider!(zxwd, X(Parameter(Val(:PiUnit), 0)), [locd])
+
+    add_edge!(zxwd, loca, locc)
+    add_edge!(zxwd, locb, locd)
+
+    locm = add_spider!(zxwd, X(Parameter(Val(:PiUnit), 0)), [loca, locb])
+    locm = add_spider!(zxwd, D, [locm])
+    locm = add_spider!(zxwd, X(Parameter(Val(:PiUnit), 1.0)), [locm])
+    add_spider!(zxwd, Z(Parameter(Val(:PiUnit), 0)), [locm, locc, locd])
+
+    # pink spider is different from red spider, we had three of them
+    # each with three legs, 3 * (3-2)/2 powers of 2 need to be added
+    # see 2307.01803
+    add_power!(zxwd,3)
+    return zxwd
+end
+
+"""
+Prepare spider at loc for integration.
+
+Perform the simplified step of zeroing out phase of spider
+and readying it for integration
+1. If target spider is X spider, turn it to Z by adding H to all its legs
+2. Pull out the Phase of the spider
+3. zero out the phase
+4. change the current spider back to its original type if necessary,
+ will generate one extra H spider.
+"""
+function int_prep!(zxwd::ZXWDiagram{T,P}, loc::T) where {T,P}
+    set_phase!(zxwd, loc, Parameter(Val(:PiUnit), 0.0))
+
+    new_loc = @match spider_type(zxwd, loc) begin
+        X(_) => add_spider!(zxwd, H, [loc])
+        Z(_) => loc
+        _ => error("Not a valid Spider to integrate over")
+    end
+    return new_loc
+end

src/parameter.jl

@@ -42,7 +42,7 @@
     Factor(f, _) => Parameter(Val(:Factor), f)
 end
 
-function Base.show(io::IO, p::Parameter)
+function Base.show(io::IO, ::MIME"text/plain", p::Parameter)
     @match p begin
         PiUnit(pu, _) && if pu isa Number
         end => print(io, "$(pu)⋅π")
@@ -70,9 +70,27 @@
 Base.:(==)(p1::Parameter, p2::Number) = eqeq(p1, p2)
 Base.:(==)(p1::Number, p2::Parameter) = eqeq(p2, p1)
 
-# following the same convention in Phase.jl implementation
+function Base.contains(p::Parameter, θ::Symbol)
+    @match p begin
+        PiUnit(pu, pt) && if !(pt <: Number)
+        end => Base.contains(repr(pu), ":" * string(θ))
+        _ => false
+    end
+end
+
+function Base.contains(p::Parameter, θ::Expr)
+    @match p begin
+        PiUnit(pu, pt) && if !(pt <: Number)
+        end => Base.contains(repr(pu), ":(" * string(θ) * ")")
+        _ => false
+    end
+end
+
+# following the same convention in phase.jl implementation
 # comparison have inconsistent, we are comparing phases to numbers
 # if cause trouble, will change
+#
+
 Base.isless(p1::Parameter, p2::Number) = @match p1 begin
     PiUnit(_...) => p1.pu isa Number && p1.pu < p2
     _ => p1.f < p2
@@ -109,8 +127,9 @@
     @match (p1, p2) begin
         (PiUnit(pu1, _), PiUnit(pu2, _)) && if pu1 isa Number && pu2 isa Number
         end => Parameter(Val(:PiUnit), pu1 + pu2)
-        (PiUnit(pu1, _), PiUnit(pu2, _)) && if !(pu1 isa Number) || !(pu2 isa Number)
-        end => Parameter(Val(:PiUnit), Expr(:call, :+, pu1, pu2))
+        (PiUnit(pu1, pu1_t), PiUnit(pu2, pu2_t)) &&
+            if !(pu1 isa Number) || !(pu2 isa Number)
+            end => PiUnit(Expr(:call, :+, pu1, pu2), Base.promote_op(+, pu1_t, pu2_t))
         (Factor(f1, _), Factor(f2, _)) => Parameter(Val(:Factor), f1 + f2)
         (PiUnit(pu1, _), Factor(f2, _)) => Parameter(Val(:Factor), exp(im * pu1 * π) * f2)
         (Factor(f1, _), PiUnit(pu2, _)) => Parameter(Val(:Factor), exp(im * pu2 * π) * f1)
@@ -134,8 +153,8 @@
         end => Parameter(Val(:PiUnit), pu1 - pu2)
         (PiUnit(pu1, pu_t1), PiUnit(pu2, pu_t2)) &&
             if !(pu1 isa Number) || !(pu2 isa Number)
-            end => Parameter(Val(:PiUnit), Expr(:call, :-, pu1, pu2))
-        (Factor(f1, _), Factor(f2, _)) => Factor(f1 - f2)
+            end => PiUnit(Expr(:call, :-, pu1, pu2), Base.promote_op(-, pu_t1, pu_t2))
+        (Factor(f1, _), Factor(f2, _)) => Parameter(Val(:Factor), f1 - f2)
         (PiUnit(_...), Factor(_...)) => Parameter(Val(:Factor), exp(im * p1.pu * π) - p2.f)
         (Factor(_...), PiUnit(_...)) => Parameter(Val(:Factor), p1.f - exp(im * p2.pu * π))
         (_, PiUnit(_...)) => Parameter(Val(:PiUnit), p1 - p2.pu)
@@ -151,6 +170,29 @@
 Base.:(-)(p1::Number, p2::Parameter) = subt_param(p1, p2)
 Base.:(-)(p1::Parameter, p2::Number) = add_param(p1, -p2)
 
+# needed for tensor contraction
+function mul_param(p1, p2)
+    @match (p1, p2) begin
+        (PiUnit(p1u, _), PiUnit(p2u, _)) && if p1u isa Number && p2u isa Number
+        end => exp(im * (p1u + p2u) * π)
+        (PiUnit(p1u, _), n2::Number) && if p1u isa Number
+        end => exp(im * p1u * π) * n2
+        (Factor(f1, _), PiUnit(p2u, _)) && if p2u isa Number
+        end => f1 * exp(im * p2u * π)
+        (PiUnit(p1u, _), Factor(f2, _)) && if p1u isa Number
+        end => f2 * exp(im * p1u * π)
+        (Factor(f1, _), Factor(f2, _)) => f1 * f2
+        (Factor(f1, _), n2::Number) => f1 * n2
+        _ => error(
+            "Invalid input '$(p1)' of type $(typeof(p1)) and '$(p2)' of type $(typeof(p2)) for ADT: *",
+        )
+    end
+end
+
+Base.:(*)(p1::Parameter, p2::Parameter) = mul_param(p1, p2)
+Base.:(*)(p1::Parameter, p2::Number) = mul_param(p1, p2)
+Base.:(*)(p1::Number, p2::Parameter) = mul_param(p2, p1)
+
 function Base.rem(p::Parameter, d::Number)
     @match p begin
         PiUnit(pu, pu_t) && if pu isa Number
@@ -163,3 +205,14 @@
         )
     end
 end
+
+function Base.inv(p::Parameter)
+    @match p begin
+        PiUnit(pu, pu_t) && if pu isa Number
+        end => Parameter(Val(:PiUnit), -pu)
+        PiUnit(pu, pu_t) && if !(pu isa Number)
+        end => Parameter(Val(:PiUnit), Expr(:call, :-, pu))
+        Factor(f, _) => Parameter(Val(:Factor), inv(f))
+        _ => error("Invalid input '$(p)' of type $(typeof(p)) for ADT: inv")
+    end
+end

src/to_eincode.jl

@@ -23,18 +23,25 @@ function to_eincode(zxwd::ZXWDiagram{T,P}) where {T,P}
             Input(q) => nothing
             Output(q) => nothing
         end
-        if res !== nothing
+
+        if res != nothing
             push!(ixs, to_eincode_indices(zxwd, v))
             push!(tensors, res)
-        else
-            push!(iy, to_eincode_indices(zxwd, v)[])
         end
     end
 
+    for v in get_outputs(zxwd)
+        push!(iy, to_eincode_indices(zxwd, v)[])
+    end
+
+    for v in get_inputs(zxwd)
+        push!(iy, to_eincode_indices(zxwd, v)[])
+    end
+
     scalar_tensor = zeros(ComplexF64, ())
 
     scalar_tensor[] = ZXCalculus.unwrap_scalar(scalar(zxwd))
-    push!(ixs, [])
+    push!(ixs, Tuple{T,T,T}[])
     push!(tensors, scalar_tensor)
     return EinCode(ixs, iy), tensors
 end
@@ -60,6 +67,7 @@ function to_eincode_indices(zxwd::ZXWDiagram{T,P}, v) where {T,P}
     end
     return ids
 end
+
 edge_index(v1, v2, mul) = (min(v1, v2), max(v1, v2), mul)
 
 function z_tensor(n::Int, α::Parameter)
@@ -81,7 +89,7 @@ function x_tensor(n::Int, α::Parameter)
     shape = (fill(2, n)...,)
     factor = @match α begin
         PiUnit(pu, _) => exp(im * pu * π)
-        Factor(f,_) => f
+        Factor(f, _) => f
         _ => error("Invalid parameter type for X-spider")
     end
     return reshape(reduce(kron, fill(pos, n)) + factor * reduce(kron, fill(neg, n)), shape)