Make the Scalar class compatible with symbolic phases

pyzx/graph/jsonparser.py

@@ -35,7 +35,7 @@
     from .multigraph import Multigraph
 
 
-def string_to_phase(string: str, g: Union[BaseGraph,'GraphDiff']) -> Union[Fraction, Poly]:
+def string_to_phase(string: str, g: Optional[Union[BaseGraph,'GraphDiff']] = None) -> Union[Fraction, Poly]:
     if not string:
         return Fraction(0)
     try:
@@ -57,6 +57,8 @@ def string_to_phase(string: str, g: Union[BaseGraph,'GraphDiff']) -> Union[Fract
             return Fraction(int(s))
     except ValueError:
         def _new_var(name: str) -> Poly:
+            if g is None:
+                return new_var(name, is_bool=False)
             return new_var(name, is_bool=g.var_registry.get_type(name, False), registry=g.var_registry)
         try:
             return parse(string, _new_var)

pyzx/graph/scalar.py

@@ -23,7 +23,7 @@
 from typing import Dict, List, Any, Union
 import json
 
-from ..utils import FloatInt, FractionLike
+from ..utils import FractionLike, phase_is_pauli, phase_is_clifford
 from ..symbolic import Poly
 
 __all__ = ['Scalar']
@@ -56,6 +56,7 @@ def __init__(self) -> None:
         self.power2: int = 0 # Stores power of square root of two
         self.phase: FractionLike = Fraction(0) # Stores complex phase of the number
         self.phasenodes: List[FractionLike] = [] # Stores list of legless spiders, by their phases.
+        self.sum_of_phases: Dict[FractionLike, int] = {} # Represents the term (c1*exp(i*phase1) + ... + cn*exp(i*phaseN)). The dictionary maps phase -> coefficient.
         self.floatfactor: complex = 1.0
         self.is_unknown: bool = False # Whether this represents an unknown scalar value
         self.is_zero: bool = False
@@ -64,16 +65,20 @@ def __repr__(self) -> str:
         return "Scalar({})".format(str(self))
 
     def __str__(self) -> str:
-        if self.is_unknown:
-            return "UNKNOWN"
-        s = "{0.real:.2f}{0.imag:+.2f}i = ".format(self.to_number())
+        s = ""
+        try:
+            s += "{0.real:.2f}{0.imag:+.2f}i = ".format(self.to_number())
+        except:
+            pass
         if self.floatfactor != 1.0:
             s += "{0.real:.2f}{0.imag:+.2f}i".format(self.floatfactor)
-        if self.phase:
-            s += "exp({}ipi)".format(str(self.phase))
         s += "sqrt(2)^{:d}".format(self.power2)
+        if self.phase:
+            s += phase_to_str(self.phase)
         for node in self.phasenodes:
-            s += "(1+exp({}ipi))".format(str(node))
+            s += "(1+exp(({})ipi))".format(str(node))
+        if self.sum_of_phases:
+            s += "(" + " + ".join([f"{coeff_to_str(coeff)}{phase_to_str(phase)}" for phase, coeff in self.sum_of_phases.items()]) + ")"
         return s
 
     def __complex__(self) -> complex:
@@ -103,8 +108,12 @@ def copy(self, conjugate: bool = False) -> 'Scalar':
         s.floatfactor = self.floatfactor if not conjugate else self.floatfactor.conjugate()
         s.is_unknown = self.is_unknown
         s.is_zero = self.is_zero
+        if not conjugate:
+            s.sum_of_phases = copy.deepcopy(self.sum_of_phases)
+        else:
+            s.sum_of_phases = {-phase: coeff for phase, coeff in self.sum_of_phases.items()}
         return s
-    
+
     def conjugate(self) -> 'Scalar':
         """Returns a new Scalar equal to the complex conjugate"""
         return self.copy(conjugate=True)
@@ -114,6 +123,11 @@ def to_number(self) -> complex:
         val = cexp(self.phase)
         for node in self.phasenodes: # Node should be a Fraction
             val *= 1+cexp(node)
+        sum_of_phases_val = 0j
+        for phase, coeff in self.sum_of_phases.items():
+            sum_of_phases_val += coeff * cexp(phase)
+        if sum_of_phases_val != 0:
+            val *= sum_of_phases_val
         val *= math.sqrt(2)**self.power2
         return val*self.floatfactor
 
@@ -136,9 +150,23 @@ def to_latex(self) -> str:
             s += r"\sqrt{{2}}^{{{:d}}}".format(self.power2)
         if self.phase not in (0,1):
             if isinstance(self.phase, Poly):
-                s += fr"\exp(i~{str(self.phase)})".format(str(self.phase))
+                s += fr"\exp(i\pi ~({str(self.phase)}))"
             else:
                 s += r"\exp(i~\frac{{{:d}\pi}}{{{:d}}})".format(self.phase.numerator,self.phase.denominator)
+
+        if self.sum_of_phases:
+            terms = []
+            for phase, coeff in self.sum_of_phases.items():
+                coeff_str = coeff_to_str(coeff)
+                if isinstance(phase, Poly):
+                    phase_str = fr"\exp(i\pi ~({str(phase)}))"
+                else:
+                    phase_str = r"\exp(i~\frac{{{:d}\pi}}{{{:d}}})".format(
+                        phase.numerator, phase.denominator)
+                terms.append(f"{coeff_str}{phase_str}")
+            if terms:
+                s += "(" + " + ".join(terms) + ")"
+
         s += "$"
         if s == "$$": return ""
         return s
@@ -151,17 +179,17 @@ def to_unicode(self) -> str:
         f = self.floatfactor
         for node in self.phasenodes:
             f *= 1+cexp(node)
-        if isinstance(self.phase, Poly):
-            raise NotImplementedError("Unicode representation of Poly not implemented")
-        phase = Fraction(self.phase)
-        if self.phase >= 1:
+        s = ""
+        phase = self.phase
+        if not isinstance(phase, Poly):
+            phase = Fraction(phase)
+        if phase >= 1:
             f *= -1
             phase -= 1
 
         if abs(f+1) > 0.001 and abs(f-1) > 0.001:
             return str(f)
 
-        s = ""
         if abs(f+1) < 0.001: #f \approx -1
             s += "-"
         if self.power2 != 0:
@@ -171,11 +199,25 @@ def to_unicode(self) -> str:
             val = str(abs(self.power2))
             s += "".join([unicode_superscript[i] for i in val])
         if phase != 0:
-            s += "exp(i"
-            if phase in unicode_fractions:
-                s += unicode_fractions[phase] + "π)"
+            if isinstance(phase, Fraction) and phase in unicode_fractions:
+                s += "exp(i" + unicode_fractions[phase] + "π)"
             else:
-                s += "{:d}/{:d}π)".format(phase.numerator,phase.denominator)
+                s += phase_to_str(phase, unicode=True)
+        if self.sum_of_phases:
+            terms = []
+            for ph, coeff in self.sum_of_phases.items():
+                term = coeff_to_str(coeff)
+                if isinstance(ph, Fraction) and ph in unicode_fractions:
+                    term += "exp(i" + unicode_fractions[ph] + "π)"
+                else:
+                    term += phase_to_str(phase, unicode=True)
+                terms.append(term)
+            if len(terms) == 1:
+                s += terms[0]
+            elif len(terms) > 1:
+                s += "(" + " + ".join(terms) + ")"
+        if s == "":
+            s = "1"
         return s
 
     def to_dict(self) -> Dict[str, Any]:
@@ -184,6 +226,8 @@ def to_dict(self) -> Dict[str, Any]:
             d["floatfactor"] =  str(self.floatfactor)
         if self.phasenodes:
             d["phasenodes"] = [str(p) for p in self.phasenodes]
+        if self.sum_of_phases:
+            d["sum_of_phases"] = {str(phase): coeff for phase, coeff in self.sum_of_phases.items()}
         if self.is_zero:
             d["is_zero"] = self.is_zero
         if self.is_unknown:
@@ -195,18 +239,21 @@ def to_json(self) -> str:
 
     @classmethod
     def from_json(cls, s: Union[str,Dict[str,Any]]) -> 'Scalar':
+        from .jsonparser import string_to_phase
+
         if isinstance(s, str):
             d = json.loads(s)
         else:
             d = s
-        # print('scalar from json', repr(d))
         scalar = Scalar()
         scalar.phase = Fraction(d["phase"]) # TODO support parameters
         scalar.power2 = int(d["power2"])
         if "floatfactor" in d:
             scalar.floatfactor = complex(d["floatfactor"])
         if "phasenodes" in d:
             scalar.phasenodes = [Fraction(p) for p in d["phasenodes"]]
+        if "sum_of_phases" in d:
+            scalar.sum_of_phases = {string_to_phase(phase): coeff for phase, coeff in d["sum_of_phases"].items()}
         if "is_zero" in d:
             scalar.is_zero = bool(d["is_zero"])
         if "is_unknown" in d:
@@ -215,14 +262,15 @@ def from_json(cls, s: Union[str,Dict[str,Any]]) -> 'Scalar':
 
     def set_unknown(self) -> None:
         self.is_unknown = True
-        self.phasenodes = []
 
     def add_power(self, n) -> None:
         """Adds a factor of sqrt(2)^n to the scalar."""
         self.power2 += n
+
     def add_phase(self, phase: FractionLike) -> None:
         """Multiplies the scalar by a complex phase."""
         self.phase = (self.phase + phase) % 2
+
     def add_node(self, node: FractionLike) -> None:
         """A solitary spider with a phase ``node`` is converted into the
         scalar 1+e^(i*pi*node)."""
@@ -231,16 +279,33 @@ def add_node(self, node: FractionLike) -> None:
         else:
             self.phasenodes.append(node)
         if node == 1: self.is_zero = True
+
     def add_float(self,f: complex) -> None:
         if f == 0.0:
             self.is_zero = True
         self.floatfactor *= f
 
+    def multiply_sum_of_phases(self, phases: Dict[FractionLike,int]) -> None:
+        """Adds a sum of phases to the scalar. The input is a dictionary mapping
+        phase -> coefficient in the sum."""
+        new_sum_of_phases: Dict[FractionLike, int] = {}
+        # Use the distributive law: (a*e^ip1)(b*e^ip2) = (a*b)*e^i(p1+p2)
+        if not self.sum_of_phases:
+            self.sum_of_phases = {0:1}
+        for phase1, coeff1 in phases.items():
+            for phase2, coeff2 in self.sum_of_phases.items():
+                new_phase = phase1 + phase2
+                new_coeff = coeff1 * coeff2
+                # Add the resulting term, combining with any existing term that has the same new polynomial
+                new_sum_of_phases[new_phase] = new_sum_of_phases.get(new_phase, 0) + new_coeff
+        self.sum_of_phases = new_sum_of_phases
+
     def mult_with_scalar(self, other: 'Scalar') -> None:
         """Multiplies two instances of Scalar together."""
         self.power2 += other.power2
         self.phase = (self.phase +other.phase)%2
         self.phasenodes.extend(other.phasenodes)
+        self.multiply_sum_of_phases(other.sum_of_phases)
         self.floatfactor *= other.floatfactor
         if other.is_zero: self.is_zero = True
         if other.is_unknown: self.is_unknown = True
@@ -250,37 +315,53 @@ def add_spider_pair(self, p1: FractionLike,p2: FractionLike) -> None:
         # These if statements look quite arbitrary, but they are just calculations of the scalar
         # of a pair of connected single wire spiders of opposite colors.
         # We make special cases for Clifford phases and pi/4 phases.
-        if p2 in (0,1):
+        if phase_is_pauli(p2):
             p1,p2 = p2, p1
-        if p1 == 0:
-            self.add_power(1)
-            return
-        elif p1 == 1:
+        if phase_is_pauli(p1):
             self.add_power(1)
-            self.add_phase(p2)
-            return
-        if isinstance(p1, Poly) or isinstance(p2, Poly):
-            self.set_unknown()
+            self.add_phase(p1 * p2)
             return
-        if p2.denominator == 2:
+
+        if phase_is_clifford(p2):
             p1, p2 = p2, p1
-        if p1 == Fraction(1,2):
-            self.add_phase(Fraction(1,4))
-            self.add_node((p2-Fraction(1,2))%2)
+        if phase_is_clifford(p1) and not phase_is_pauli(p1):
+            self.add_phase(Fraction(3,2) * p1 - Fraction(1,2))
+            self.add_node((p2-p1)%2)
             return
-        elif p1 == Fraction(3,2):
-            self.add_phase(Fraction(7,4))
-            self.add_node((p2-Fraction(3,2))%2)
-            return
-        if (p1 + p2) % 2 == 0:
+
+        if (p1 + p2) % 2 == 0 and not isinstance(p1, Poly) and not isinstance(p2, Poly):
             if p1.denominator == 4:
                 if p1.numerator in (3,5):
                     self.add_phase(Fraction(1))
                 return
             self.add_power(1)
             self.add_float(math.cos(p1))
             return
+
         # Generic case
         self.add_power(-1)
-        self.add_float(1+cexp(p1)+cexp(p2) - cexp(p1+p2))
+        # add the sum of phases 1 + e^(i pi p1) + e^(i pi p2) - e^(i pi (p1+p2))
+        self.multiply_sum_of_phases({0:1, p1:1, p2:1, (p1+p2)%2:-1})
         return
+
+
+def coeff_to_str(coeff: int) -> str:
+    if coeff == 1:
+        return ""
+    if coeff == -1:
+        return "-"
+    return str(coeff) + "*"
+
+def phase_to_str(phase: FractionLike, unicode: bool = False) -> str:
+    pi = "π" if unicode else "pi"
+    if phase == 0:
+        return "1"
+    if phase == 1:
+        return f"exp(i{pi})"
+    if isinstance(phase, Poly):
+        if len(phase.terms) == 1:
+            return f"exp({str(phase)}i{pi})"
+        return f"exp(({str(phase)})i{pi})"
+    if isinstance(phase, Fraction):
+        return f"exp({phase.numerator}/{phase.denominator}i{pi})"
+    return f"exp({str(phase)}i{pi})"