several fixes and improvements to multigraph

pyzx/graph/diff.py

@@ -84,13 +84,13 @@ def calculate_diff(self, g1: BaseGraph[VT,ET], g2: BaseGraph[VT,ET]) -> None:
 				pos2 = g2.qubit(v), g2.row(v)
 				self.changed_pos[v] = pos2
 
-		# for e in new_edges:
-		# 	if e in old_edges:
-		# 		if g1.edge_type(e) != g2.edge_type(e):
-		# 			self.changed_edge_types[e] = g2.edge_type(e)
-		# 	else:
-		# 		if g2.edge_type(e) != EdgeType.HADAMARD: # We take Hadamard edges to be the default
-		# 			self.changed_edge_types[e] = g2.edge_type(e)
+		for e in new_edges:
+			if e in old_edges:
+				if g1.edge_type(e) != g2.edge_type(e):
+					self.changed_edge_types[e] = g2.edge_type(e)
+			else:
+				if g2.edge_type(e) != EdgeType.HADAMARD: # We take Hadamard edges to be the default
+					self.changed_edge_types[e] = g2.edge_type(e)
 
 	def apply_diff(self,g: BaseGraph[VT,ET]) -> BaseGraph[VT,ET]:
 		g = copy.deepcopy(g)

pyzx/graph/multigraph.py

@@ -14,6 +14,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections import Counter
 from fractions import Fraction
 from typing import Tuple, Dict, Set, Any
 
@@ -26,21 +27,34 @@ class Edge:
     between two vertices"""
     s: int
     h: int
-    def __init__(self, s: int=0, h: int=0):
+    w_io: int
+
+    def __init__(self, s: int=0, h: int=0, w_io: int=0):
         self.s = s
         self.h = h
+        self.w_io = w_io
 
-    def add(self, s: int=0, h: int=0):
+    def add(self, s: int=0, h: int=0, w_io: int=0):
         self.s += s
         self.h += h
+        self.w_io += w_io
         if self.s < 0 or self.h < 0:
             raise ValueError('Cannot have negative edges')
+        if self.w_io not in (0,1):
+            raise ValueError('Invalid number of W-IO edges')
+        if self.w_io == 1 and self.s + self.h > 0:
+            raise ValueError('Cannot have W-IO edge and other edges')
+
+    def remove(self, s: int=0, h: int=0, w_io: int=0):
+        self.add(s=-s, h=-h, w_io=-w_io)
 
-    def remove(self, s: int=0, h: int=0):
-        self.add(s=-s, h=-h)
-
     def is_empty(self) -> bool:
-        return self.s == 0 and self.h == 0
+        return self.s == 0 and self.h == 0 and self.w_io == 0
+
+    def get_edge_count(self, ty: EdgeType.Type) -> int:
+        if ty == EdgeType.SIMPLE: return self.s
+        elif ty == EdgeType.HADAMARD: return self.h
+        else: return self.w_io
 
 class Multigraph(BaseGraph[int,Tuple[int,int,EdgeType.Type]]):
     """Purely Pythonic multigraph implementation of :class:`~graph.base.BaseGraph`."""
@@ -156,7 +170,8 @@ def add_edge(self, edge_pair, edgetype=EdgeType.SIMPLE):
             e = self.graph[s][t]
 
         if edgetype == EdgeType.SIMPLE: e.add(s=1)
-        else: e.add(h=1)
+        elif edgetype == EdgeType.HADAMARD: e.add(h=1)
+        else: e.add(w_io=1)
 
         if self._auto_simplify:
             t1 = self.ty[s]
@@ -224,9 +239,10 @@ def remove_edge(self, edge):
         s,t,ty = edge
         e = self.graph[s][t]
         if ty == EdgeType.SIMPLE: e.remove(s=1)
-        else: e.remove(h=1)
+        elif ty == EdgeType.HADAMARD: e.remove(h=1)
+        else: e.remove(w_io=1)
 
-        if e.s == 0 and e.h == 0:
+        if e.is_empty():
             del self.graph[s][t]
             del self.graph[t][s]
 
@@ -257,11 +273,13 @@ def edges(self, s=None, t=None):
                     if v1 > v0:
                         for _ in range(e.s): yield (v0, v1, EdgeType.SIMPLE)
                         for _ in range(e.h): yield (v0, v1, EdgeType.HADAMARD)
+                        for _ in range(e.w_io): yield (v0, v1, EdgeType.W_IO)
         elif t != None:
             s, t = (s, t) if s < t else (t, s)
             e = self.graph[s][t]
             for _ in range(e.s): yield (s, t, EdgeType.SIMPLE)
             for _ in range(e.h): yield (s, t, EdgeType.HADAMARD)
+            for _ in range(e.w_io): yield (s, t, EdgeType.W_IO)
 
     # def edges_in_range(self, start, end, safe=False):
     #    """like self.edges, but only returns edges that belong to vertices
@@ -287,8 +305,10 @@ def edges(self, s=None, t=None):
 
     def edge(self, s, t):
         return (s,t) if s < t else (t,s)
+
     def edge_set(self):
-        return set(self.edges())
+        return Counter(self.edges())
+
     def edge_st(self, edge):
         return (edge[0], edge[1])
 
@@ -309,14 +329,24 @@ def connected(self,v1,v2):
         return v2 in self.graph[v1]
 
     def edge_type(self, e):
+        if len(e) == 2:
+            edges = list(self.edges(e[0],e[1]))
+            if len(edges) != 1:
+                raise ValueError('Cannot determine edge type')
+            e = edges[0]
         return e[2]
 
     def set_edge_type(self, edge, t):
         v1,v2,ty = edge
         if ty != t:
             e = self.graph[v1][v2]
-            if t == EdgeType.SIMPLE: e.add(s=1,h=-1)
-            else: e.add(s=-1,h=1)
+            # decrement the old type and increment the new type
+            if ty == EdgeType.SIMPLE: e.add(s=-1)
+            elif ty == EdgeType.HADAMARD: e.add(h=-1)
+            else: e.add(w_io=-1)
+            if t == EdgeType.SIMPLE: e.add(s=1)
+            elif t == EdgeType.HADAMARD: e.add(h=1)
+            else: e.add(w_io=1)
 
     def type(self, vertex):
         return self.ty[vertex]