Added color to register net plot depending on the fidelity of entanglement + Added directional metadata

ext/QuantumSavoryMakie/QuantumSavoryMakie.jl

@@ -21,6 +21,10 @@ function Makie.plot!(rn::RegisterNetPlot{<:Tuple{RegisterNet}}) # TODO plot the
     state_coords = Makie.Observable(Makie.Point2{Float64}[])
     state_coords_backref = Makie.Observable([])
     state_links = Makie.Observable(Makie.Point2{Float64}[])
+    twoqubitstate_links = Makie.Observable(Makie.Point2{Float64}[])
+    clrs = Makie.Observable(Float32[])
+    regs = Makie.Observable([])
+    twoqubitobservable = rn[:twoqubitobservable][]
     if haskey(rn, :registercoords) && !isnothing(rn[:registercoords][])
         registercoords = rn[:registercoords][]
     else
@@ -32,7 +36,8 @@ function Makie.plot!(rn::RegisterNetPlot{<:Tuple{RegisterNet}}) # TODO plot the
         registers = network.registers
         all_nodes = [ # TODO it is rather wasteful to replot everything... do it smarter
             register_rectangles, register_slots_coords, register_slots_coords_backref,
-            state_coords, state_coords_backref, state_links,
+            state_coords, state_coords_backref, state_links, twoqubitstate_links,
+            clrs, regs
         ]
         for a in all_nodes
             empty!(a[])
@@ -57,13 +62,29 @@ function Makie.plot!(rn::RegisterNetPlot{<:Tuple{RegisterNet}}) # TODO plot the
                     break
                 end
                 push!(state_links[], Makie.Point2{Float64}(registercoords[whichreg][1], registercoords[whichreg][2]+i-1))
+                if nsubsystems(s)==2 && !isnothing(twoqubitobservable) 
+                    push!(regs[], (whichreg, i, s))
+                    push!(twoqubitstate_links[], Makie.Point2{Float64}(registercoords[whichreg][1], registercoords[whichreg][2]+i-1))
+                end
                 if !juststarted && si<nsubsystems(s)
                     push!(state_links[], state_links[][end])
+                    if nsubsystems(s)==2 && !isnothing(twoqubitobservable)
+                        push!(regs[], regs[][end])
+                        push!(twoqubitstate_links[], twoqubitstate_links[][end])
+                    end
                 else
                     juststarted = false
                 end
             end
         end
+        if !isnothing(twoqubitobservable)
+            for i in 1:2:length(regs[])
+                s = regs[][i][3]
+                fid = real(observable([registers[regs[][i][1]][regs[][i][2]], registers[regs[][i+1][1]][regs[][i+1][2]]], twoqubitobservable))
+                push!(clrs[], fid)
+            end
+        end
+
         for a in all_nodes
             notify(a)
         end
@@ -75,12 +96,16 @@ function Makie.plot!(rn::RegisterNetPlot{<:Tuple{RegisterNet}}) # TODO plot the
     register_slots_scatterplot = Makie.scatter!(rn,register_slots_coords,marker=:rect,color=:gray60,markersize=0.6,markerspace=:data)
     state_scatterplot = Makie.scatter!(rn,state_coords,marker=:diamond,color=:black,markersize=0.4,markerspace=:data)
     state_lineplot = Makie.linesegments!(rn,state_links,color=:gray90)
+    twoqubitstate_lineplot = Makie.linesegments!(rn,twoqubitstate_links,color= !isnothing(twoqubitobservable) ? clrs : (:gray90), colormap=:Spectral, colorrange = (0,1))
     rn[:register_polyplot] = register_polyplot
     rn[:register_slots_scatterplot] = register_slots_scatterplot
     rn[:register_slots_coords_backref] = register_slots_coords_backref
     rn[:state_scatterplot] = state_scatterplot
     rn[:state_coords_backref] = state_coords_backref
     rn[:state_lineplot] = state_lineplot
+    rn[:twoqubitstate_lineplot] = twoqubitstate_lineplot
+    rn[:twoqubitobservable] = twoqubitobservable
+    rn[:fids] = clrs
     rn
 end
 
@@ -111,10 +136,10 @@ end
 It returns a tuple of (subfigure, axis, plot, observable).
 The observable can be used to issue a `notify` call that updates
 the plot with the current state of the network."""
-function registernetplot_axis(subfig, registersobservable; registercoords=nothing, interactions=false)
+function registernetplot_axis(subfig, registersobservable; registercoords=nothing, interactions=false, twoqubitobservable=nothing)
     ax = Makie.Axis(subfig)
     p = registernetplot!(ax, registersobservable,
-        registercoords=registercoords,
+        registercoords=registercoords, twoqubitobservable=twoqubitobservable
         )
     ax.aspect = Makie.DataAspect()
     Makie.hidedecorations!(ax)

src/QuantumSavory.jl

@@ -107,7 +107,8 @@ struct RegisterNet
     registers::Vector{Register}
     vertex_metadata::Vector{Dict{Symbol,Any}}
     edge_metadata::Dict{Tuple{Int,Int},Dict{Symbol,Any}}
-    function RegisterNet(graph, registers, vertex_metadata, edge_metadata)
+    directed_edge_metadata::Dict{Pair{Int,Int},Dict{Symbol,Any}}
+    function RegisterNet(graph, registers, vertex_metadata, edge_metadata, directed_edge_metadata)
         all_are_at_zero = all(iszero(ConcurrentSim.now(r.env)) && isempty(r.env.heap) && isnothing(r.env.active_proc) for r in registers)
         all_are_same = all(registers[1].env === r.env for r in registers)
         if !all_are_same
@@ -123,7 +124,7 @@ struct RegisterNet
                 error("When constructing a `RegisterNet`, the registers must either have not been used yet or have to already belong to the same simulation time tracker, which is not the case here. The simplest way to fix this error is to immediately construct the `RegisterNet` after you have constructed the registers.")
             end
         end
-        new(graph, registers, vertex_metadata, edge_metadata)
+        new(graph, registers, vertex_metadata, edge_metadata, directed_edge_metadata)
     end
 end
 """
@@ -162,7 +163,7 @@ julia> neighbors(net, 1)
 """
 function RegisterNet(graph::SimpleGraph, registers::Vector{Register})
     @assert size(graph, 1) == length(registers)
-    RegisterNet(graph, registers, [Dict{Symbol,Any}() for _ in registers], Dict{Tuple{Int,Int},Dict{Symbol,Any}}())
+    RegisterNet(graph, registers, [Dict{Symbol,Any}() for _ in registers], Dict{Tuple{Int,Int},Dict{Symbol,Any}}(), Dict{Pair{Int,Int},Dict{Symbol,Any}}())
 end
 """Construct a [`RegisterNet`](@ref) from a given list of [`Register`](@ref)s, defaulting to a chain topology.
 
@@ -209,13 +210,22 @@ function Base.setindex!(net::RegisterNet, val, ij::Tuple{Int,Int}, k::Symbol)
     haskey(net.edge_metadata,edge) || (net.edge_metadata[edge] = Dict{Symbol,Any}())
     net.edge_metadata[edge][k] = val
 end
+# Get and set directed edge metadata
+Base.getindex(net::RegisterNet, ij::Pair{Int,Int}, k::Symbol) = net.directed_edge_metadata[ij][k]
+function Base.setindex!(net::RegisterNet, val, ij::Pair{Int,Int}, k::Symbol)
+    edge = ij
+    haskey(net.directed_edge_metadata,edge) || (net.directed_edge_metadata[edge] = Dict{Symbol,Any}())
+    net.directed_edge_metadata[edge][k] = val
+end
 Base.getindex(net::RegisterNet, ij::Graphs.SimpleEdge, k::Symbol) = net[(ij.src, ij.dst),k]
 Base.setindex!(net::RegisterNet, val, ij::Graphs.SimpleEdge, k::Symbol) = begin net[(ij.src, ij.dst),k] = val end
 # Get and set with colon notation
 Base.getindex(net::RegisterNet, ::Colon) = net.registers
 Base.getindex(net::RegisterNet, ::Colon, j::Int) = [r[j] for r in net.registers]
 Base.getindex(net::RegisterNet, ::Colon, k::Symbol) = [m[k] for m in net.vertex_metadata]
 Base.getindex(net::RegisterNet, ::Tuple{Colon,Colon}, k::Symbol) = [net.edge_metadata[minmax(ij)...][k] for ij in edges(net)]
+Base.getindex(net::RegisterNet, ::Pair{Colon,Colon}, k::Symbol) = [net.directed_edge_metadata[ij][k] for ij in edges(net)]
+
 function Base.setindex!(net::RegisterNet, v, ::Colon, k::Symbol)
     for m in net.vertex_metadata
         m[k] = v
@@ -226,6 +236,11 @@ function Base.setindex!(net::RegisterNet, v, ::Tuple{Colon,Colon}, k::Symbol)
         net[ij,k] = v
     end
 end
+function Base.setindex!(net::RegisterNet, v, ::Pair{Colon,Colon}, k::Symbol)
+    for ij in edges(net)
+        net[ij,k] = v
+    end
+end
 function Base.setindex!(net::RegisterNet, @nospecialize(f::Function), ::Colon, k::Symbol)
     for m in net.vertex_metadata
         m[k] = f()
@@ -236,6 +251,11 @@ function Base.setindex!(net::RegisterNet, @nospecialize(f::Function), ::Tuple{Co
         net[ij,k] = f()
     end
 end
+function Base.setindex!(net::RegisterNet, @nospecialize(f::Function), ::Pair{Colon,Colon}, k::Symbol)
+    for ij in edges(net)
+        net[ij,k] = f()
+    end
+end
 
 ##