Allow symbolic scalars in linear combinations (#7030)

* Parameterize LinearDict

* Add protocol handlers

* fix protocol handlers

* fix protocol handlers, add test

* format

* tests

* Y, Z gates, parameterize test

* mypy

* Revert changes to Scalar, and just use TParamValComplex everywhere. Add tests.

* Parameterize LinearCombinationOfGates

* tests

* tests

* tests

* tests

* test

* Update cirq-core/cirq/ops/raw_types.py

Co-authored-by: Pavol Juhas <pavol.juhas@gmail.com>

---------

Co-authored-by: Pavol Juhas <pavol.juhas@gmail.com>

Author: daxfohl

cirq-core/cirq/linalg/operator_spaces.py

@@ -13,13 +13,17 @@
 # limitations under the License.
 
 """Utilities for manipulating linear operators as elements of vector space."""
-from typing import Dict, Tuple
+from typing import Dict, Tuple, TYPE_CHECKING
 
 import numpy as np
+import sympy
 
 from cirq import value
 from cirq._doc import document
 
+if TYPE_CHECKING:
+    import cirq
+
 PAULI_BASIS = {
     'I': np.eye(2),
     'X': np.array([[0.0, 1.0], [1.0, 0.0]]),
@@ -78,8 +82,17 @@ def matrix_from_basis_coefficients(
 
 
 def pow_pauli_combination(
-    ai: value.Scalar, ax: value.Scalar, ay: value.Scalar, az: value.Scalar, exponent: int
-) -> Tuple[value.Scalar, value.Scalar, value.Scalar, value.Scalar]:
+    ai: 'cirq.TParamValComplex',
+    ax: 'cirq.TParamValComplex',
+    ay: 'cirq.TParamValComplex',
+    az: 'cirq.TParamValComplex',
+    exponent: int,
+) -> Tuple[
+    'cirq.TParamValComplex',
+    'cirq.TParamValComplex',
+    'cirq.TParamValComplex',
+    'cirq.TParamValComplex',
+]:
     """Computes non-negative integer power of single-qubit Pauli combination.
 
     Returns scalar coefficients bi, bx, by, bz such that
@@ -96,7 +109,10 @@ def pow_pauli_combination(
     if exponent == 0:
         return 1, 0, 0, 0
 
-    v = np.sqrt(ax * ax + ay * ay + az * az).item()
+    if any(isinstance(a, sympy.Basic) for a in [ax, ay, az]):
+        v = sympy.sqrt(ax * ax + ay * ay + az * az)
+    else:
+        v = np.sqrt(ax * ax + ay * ay + az * az).item()
     s = (ai + v) ** exponent
     t = (ai - v) ** exponent
 

cirq-core/cirq/linalg/operator_spaces_test.py

@@ -17,6 +17,7 @@
 import numpy as np
 import pytest
 import scipy.linalg
+import sympy
 
 import cirq
 
@@ -287,11 +288,21 @@ def test_expand_is_inverse_of_reconstruct(m1, basis):
             (-1, -2, 3, 4),
             (1j, 2j, 3j, 4j),
             (1j, 2j, 3, 4),
+            (sympy.Symbol('i'), sympy.Symbol('x'), sympy.Symbol('y'), sympy.Symbol('z')),
+            (
+                sympy.Symbol('i') * 1j,
+                -sympy.Symbol('x'),
+                -sympy.Symbol('y') * 1j,
+                sympy.Symbol('z'),
+            ),
         ),
         (0, 1, 2, 3, 4, 5, 100, 101),
     ),
 )
 def test_pow_pauli_combination(coefficients, exponent):
+    is_symbolic = any(isinstance(a, sympy.Basic) for a in coefficients)
+    if is_symbolic and exponent > 2:
+        return  # too slow
     i = cirq.PAULI_BASIS['I']
     x = cirq.PAULI_BASIS['X']
     y = cirq.PAULI_BASIS['Y']
@@ -303,5 +314,7 @@ def test_pow_pauli_combination(coefficients, exponent):
 
     bi, bx, by, bz = cirq.pow_pauli_combination(ai, ax, ay, az, exponent)
     result = bi * i + bx * x + by * y + bz * z
-
-    assert np.allclose(result, expected_result)
+    if is_symbolic:
+        assert cirq.approx_eq(result, expected_result)
+    else:
+        assert np.allclose(result, expected_result)

cirq-core/cirq/ops/common_gates.py

@@ -261,8 +261,8 @@ def _pauli_expansion_(self) -> value.LinearDict[str]:
         if self._dimension != 2:
             return NotImplemented
         phase = 1j ** (2 * self._exponent * (self._global_shift + 0.5))
-        angle = _pi(self._exponent) * self._exponent / 2
         lib = sympy if protocols.is_parameterized(self) else np
+        angle = lib.pi * self._exponent / 2
         return value.LinearDict({'I': phase * lib.cos(angle), 'X': -1j * phase * lib.sin(angle)})
 
     def _circuit_diagram_info_(
@@ -466,8 +466,8 @@ def _trace_distance_bound_(self) -> Optional[float]:
 
     def _pauli_expansion_(self) -> value.LinearDict[str]:
         phase = 1j ** (2 * self._exponent * (self._global_shift + 0.5))
-        angle = _pi(self._exponent) * self._exponent / 2
         lib = sympy if protocols.is_parameterized(self) else np
+        angle = lib.pi * self._exponent / 2
         return value.LinearDict({'I': phase * lib.cos(angle), 'Y': -1j * phase * lib.sin(angle)})
 
     def _circuit_diagram_info_(
@@ -767,8 +767,8 @@ def _pauli_expansion_(self) -> value.LinearDict[str]:
         if self._dimension != 2:
             return NotImplemented
         phase = 1j ** (2 * self._exponent * (self._global_shift + 0.5))
-        angle = _pi(self._exponent) * self._exponent / 2
         lib = sympy if protocols.is_parameterized(self) else np
+        angle = lib.pi * self._exponent / 2
         return value.LinearDict({'I': phase * lib.cos(angle), 'Z': -1j * phase * lib.sin(angle)})
 
     def _phase_by_(self, phase_turns: float, qubit_index: int):

cirq-core/cirq/ops/linear_combinations.py

@@ -81,7 +81,7 @@ class LinearCombinationOfGates(value.LinearDict[raw_types.Gate]):
         2 * cirq.X - 2 * cirq.Z
     """
 
-    def __init__(self, terms: Mapping[raw_types.Gate, value.Scalar]) -> None:
+    def __init__(self, terms: Mapping[raw_types.Gate, 'cirq.TParamValComplex']) -> None:
         """Initializes linear combination from a collection of terms.
 
         Args:
@@ -149,17 +149,19 @@ def __pow__(self, exponent: int) -> 'LinearCombinationOfGates':
         )
 
     def _is_parameterized_(self) -> bool:
-        return any(protocols.is_parameterized(gate) for gate in self.keys())
+        return any(protocols.is_parameterized(item) for item in self.items())
 
     def _parameter_names_(self) -> AbstractSet[str]:
-        return {name for gate in self.keys() for name in protocols.parameter_names(gate)}
+        return {name for item in self.items() for name in protocols.parameter_names(item)}
 
     def _resolve_parameters_(
         self, resolver: 'cirq.ParamResolver', recursive: bool
     ) -> 'LinearCombinationOfGates':
         return self.__class__(
             {
-                protocols.resolve_parameters(gate, resolver, recursive): coeff
+                protocols.resolve_parameters(
+                    gate, resolver, recursive
+                ): protocols.resolve_parameters(coeff, resolver, recursive)
                 for gate, coeff in self.items()
             }
         )
@@ -222,7 +224,7 @@ class LinearCombinationOfOperations(value.LinearDict[raw_types.Operation]):
     by the identity operator. Note that A may not be unitary or even normal.
     """
 
-    def __init__(self, terms: Mapping[raw_types.Operation, value.Scalar]) -> None:
+    def __init__(self, terms: Mapping[raw_types.Operation, 'cirq.TParamValComplex']) -> None:
         """Initializes linear combination from a collection of terms.
 
         Args:
@@ -264,17 +266,19 @@ def __pow__(self, exponent: int) -> 'LinearCombinationOfOperations':
         return LinearCombinationOfOperations({i: bi, x: bx, y: by, z: bz})
 
     def _is_parameterized_(self) -> bool:
-        return any(protocols.is_parameterized(op) for op in self.keys())
+        return any(protocols.is_parameterized(item) for item in self.items())
 
     def _parameter_names_(self) -> AbstractSet[str]:
-        return {name for op in self.keys() for name in protocols.parameter_names(op)}
+        return {name for item in self.items() for name in protocols.parameter_names(item)}
 
     def _resolve_parameters_(
         self, resolver: 'cirq.ParamResolver', recursive: bool
     ) -> 'LinearCombinationOfOperations':
         return self.__class__(
             {
-                protocols.resolve_parameters(op, resolver, recursive): coeff
+                protocols.resolve_parameters(op, resolver, recursive): protocols.resolve_parameters(
+                    coeff, resolver, recursive
+                )
                 for op, coeff in self.items()
             }
         )
@@ -353,7 +357,9 @@ def _is_linear_dict_of_unit_pauli_string(linear_dict: value.LinearDict[UnitPauli
     return True
 
 
-def _pauli_string_from_unit(unit: UnitPauliStringT, coefficient: Union[int, float, complex] = 1):
+def _pauli_string_from_unit(
+    unit: UnitPauliStringT, coefficient: Union[int, float, 'cirq.TParamValComplex'] = 1
+):
     return PauliString(qubit_pauli_map=dict(unit), coefficient=coefficient)
 
 

cirq-core/cirq/ops/linear_combinations_test.py

@@ -153,6 +153,10 @@ def test_non_unitary_linear_combination_of_gates_has_no_unitary(terms):
         ),
         ({cirq.X: 2, cirq.H: 1}, {'X': 2 + np.sqrt(0.5), 'Z': np.sqrt(0.5)}),
         ({cirq.XX: -2, cirq.YY: 3j, cirq.ZZ: 4}, {'XX': -2, 'YY': 3j, 'ZZ': 4}),
+        (
+            {cirq.X: sympy.Symbol('x'), cirq.Y: -sympy.Symbol('y')},
+            {'X': sympy.Symbol('x'), 'Y': -sympy.Symbol('y')},
+        ),
     ),
 )
 def test_linear_combination_of_gates_has_correct_pauli_expansion(terms, expected_expansion):
@@ -206,7 +210,11 @@ def test_linear_combinations_of_gates_invalid_powers(terms, exponent):
 
 @pytest.mark.parametrize(
     'terms, is_parameterized, parameter_names',
-    [({cirq.H: 1}, False, set()), ({cirq.X ** sympy.Symbol('t'): 1}, True, {'t'})],
+    [
+        ({cirq.H: 1}, False, set()),
+        ({cirq.X ** sympy.Symbol('t'): 1}, True, {'t'}),
+        ({cirq.X: sympy.Symbol('t')}, True, {'t'}),
+    ],
 )
 @pytest.mark.parametrize('resolve_fn', [cirq.resolve_parameters, cirq.resolve_parameters_once])
 def test_parameterized_linear_combination_of_gates(
@@ -225,7 +233,7 @@ def get_matrix(
         cirq.GateOperation,
         cirq.LinearCombinationOfGates,
         cirq.LinearCombinationOfOperations,
-    ]
+    ],
 ) -> np.ndarray:
     if isinstance(operator, (cirq.LinearCombinationOfGates, cirq.LinearCombinationOfOperations)):
         return operator.matrix()
@@ -243,13 +251,13 @@ def assert_linear_combinations_are_equal(
 
     actual_matrix = get_matrix(actual)
     expected_matrix = get_matrix(expected)
-    assert np.allclose(actual_matrix, expected_matrix)
+    assert cirq.approx_eq(actual_matrix, expected_matrix)
 
     actual_expansion = cirq.pauli_expansion(actual)
     expected_expansion = cirq.pauli_expansion(expected)
     assert set(actual_expansion.keys()) == set(expected_expansion.keys())
     for name in actual_expansion.keys():
-        assert abs(actual_expansion[name] - expected_expansion[name]) < 1e-12
+        assert cirq.approx_eq(actual_expansion[name], expected_expansion[name])
 
 
 @pytest.mark.parametrize(
@@ -279,6 +287,8 @@ def assert_linear_combinations_are_equal(
         ),
         ((cirq.X + cirq.Y + cirq.Z) ** 0, cirq.I),
         ((cirq.X - 1j * cirq.Y) ** 0, cirq.I),
+        (cirq.Y - sympy.Symbol('s') * cirq.Y, (1 - sympy.Symbol('s')) * cirq.Y),
+        ((cirq.X + cirq.Z) * sympy.Symbol('s') / np.sqrt(2), cirq.H * sympy.Symbol('s')),
     ),
 )
 def test_gate_expressions(expression, expected_result):
@@ -659,6 +669,10 @@ def test_non_unitary_linear_combination_of_operations_has_no_unitary(terms):
             {'IIZI': 1, 'IZII': 1, 'IZZI': -1},
         ),
         ({cirq.CNOT(q0, q1): 2, cirq.Z(q0): -1, cirq.X(q1): -1}, {'II': 1, 'ZX': -1}),
+        (
+            {cirq.X(q0): -sympy.Symbol('x'), cirq.Y(q0): sympy.Symbol('y')},
+            {'X': -sympy.Symbol('x'), 'Y': sympy.Symbol('y')},
+        ),
     ),
 )
 def test_linear_combination_of_operations_has_correct_pauli_expansion(terms, expected_expansion):
@@ -716,6 +730,7 @@ def test_linear_combinations_of_operations_invalid_powers(terms, exponent):
     [
         ({cirq.H(cirq.LineQubit(0)): 1}, False, set()),
         ({cirq.X(cirq.LineQubit(0)) ** sympy.Symbol('t'): 1}, True, {'t'}),
+        ({cirq.X(cirq.LineQubit(0)): sympy.Symbol('t')}, True, {'t'}),
     ],
 )
 @pytest.mark.parametrize('resolve_fn', [cirq.resolve_parameters, cirq.resolve_parameters_once])
@@ -788,6 +803,10 @@ def test_parameterized_linear_combination_of_ops(
             cirq.LinearCombinationOfOperations({cirq.X(q1): 2, cirq.Z(q1): 3}) ** 0,
             cirq.LinearCombinationOfOperations({cirq.I(q1): 1}),
         ),
+        (
+            cirq.LinearCombinationOfOperations({cirq.X(q0): sympy.Symbol('s')}) ** 2,
+            cirq.LinearCombinationOfOperations({cirq.I(q0): sympy.Symbol('s') ** 2}),
+        ),
     ),
 )
 def test_operation_expressions(expression, expected_result):

cirq-core/cirq/ops/raw_types.py

@@ -306,7 +306,7 @@ def on_each(self, *targets: Union[Qid, Iterable[Any]]) -> List['cirq.Operation']
         return operations
 
     def wrap_in_linear_combination(
-        self, coefficient: Union[complex, float, int] = 1
+        self, coefficient: 'cirq.TParamValComplex' = 1
     ) -> 'cirq.LinearCombinationOfGates':
         """Returns a LinearCombinationOfGates with this gate.
 

cirq-core/cirq/ops/raw_types_test.py

@@ -258,6 +258,7 @@ def _decompose_(self, qubits):
         (cirq.CZ * 1, cirq.CZ / 1),
         (-cirq.CSWAP * 1j, cirq.CSWAP / 1j),
         (cirq.TOFFOLI * 0.5, cirq.TOFFOLI / 2),
+        (-cirq.X * sympy.Symbol('s'), -sympy.Symbol('s') * cirq.X),
     ),
 )
 def test_gate_algebra(expression, expected_result):