ProjectorSum object for TFQ (quantumlib#4364)

This is a follow-up on quantumlib#4331 which introduced ProjectorString objects. The present PR aims at adding ProjectorSum objects.

cirq/__init__.py

@@ -256,6 +256,7 @@
     PhasedXZGate,
     PhaseFlipChannel,
     ProjectorString,
+    ProjectorSum,
     RandomGateChannel,
     qft,
     Qid,

cirq/json_resolver_cache.py

@@ -121,6 +121,7 @@ def two_qubit_matrix_gate(matrix):
         'PhasedXPowGate': cirq.PhasedXPowGate,
         'PhasedXZGate': cirq.PhasedXZGate,
         'ProjectorString': cirq.ProjectorString,
+        'ProjectorSum': cirq.ProjectorSum,
         'RandomGateChannel': cirq.RandomGateChannel,
         'QuantumFourierTransformGate': cirq.QuantumFourierTransformGate,
         'RepetitionsStoppingCriteria': cirq.work.RepetitionsStoppingCriteria,

cirq/ops/__init__.py

@@ -129,6 +129,7 @@
     LinearCombinationOfOperations,
     PauliSum,
     PauliSumLike,
+    ProjectorSum,
 )
 
 from cirq.ops.mixed_unitary_channel import (

cirq/ops/linear_combinations.py

@@ -14,6 +14,7 @@
 from collections import defaultdict
 from typing import (
     AbstractSet,
+    Any,
     Dict,
     Iterable,
     Mapping,
@@ -31,12 +32,14 @@
 from sympy.logic.boolalg import And, Not, Or, Xor
 from sympy.core.expr import Expr
 from sympy.core.symbol import Symbol
+from scipy.sparse import csr_matrix
 
 from cirq import linalg, protocols, qis, value
 from cirq._doc import document
 from cirq.linalg import operator_spaces
 from cirq.ops import identity, raw_types, pauli_gates, pauli_string
 from cirq.ops.pauli_string import PauliString, _validate_qubit_mapping
+from cirq.ops.projector import ProjectorString
 from cirq.value.linear_dict import _format_terms
 
 if TYPE_CHECKING:
@@ -739,3 +742,210 @@ def __format__(self, format_spec: str) -> str:
 
     def __str__(self) -> str:
         return self.__format__('.3f')
+
+
+def _projector_string_from_projector_dict(projector_dict, coefficient=1.0):
+    return ProjectorString(dict(projector_dict), coefficient)
+
+
+@value.value_equality(approximate=True)
+class ProjectorSum:
+    def __init__(
+        self, linear_dict: Optional[value.LinearDict[FrozenSet[Tuple[raw_types.Qid, int]]]] = None
+    ):
+        """Constructor for ProjectorSum
+
+        Args:
+            linear_dict: A linear dictionary from a set of tuples of (Qubit, integer) to a complex
+                number. The tuple is a projector onto the qubit and the complex number is the
+                weight of these projections.
+        """
+        self._linear_dict: value.LinearDict[FrozenSet[Tuple[raw_types.Qid, int]]] = (
+            linear_dict if linear_dict is not None else value.LinearDict({})
+        )
+
+    def _value_equality_values_(self):
+        return self._linear_dict
+
+    def _json_dict_(self) -> Dict[str, Any]:
+        linear_dict = []
+        for projector_dict, scalar in dict(self._linear_dict).items():
+            key = [[k, v] for k, v in dict(projector_dict).items()]
+            linear_dict.append([key, scalar])
+        return {
+            'cirq_type': self.__class__.__name__,
+            'linear_dict': linear_dict,
+        }
+
+    @classmethod
+    def _from_json_dict_(cls, linear_dict, **kwargs):
+        converted_dict = {}
+        for projector_string in linear_dict:
+            projector_dict = {x[0]: x[1] for x in projector_string[0]}
+            scalar = projector_string[1]
+            key = frozenset(projector_dict.items())
+            converted_dict[key] = scalar
+        return cls(linear_dict=value.LinearDict(converted_dict))
+
+    @classmethod
+    def from_projector_strings(
+        cls, terms: Union[ProjectorString, List[ProjectorString]]
+    ) -> 'ProjectorSum':
+        """Builds a ProjectorSum from one or more ProjectorString(s).
+
+        Args:
+            terms: Either a single ProjectorString or a list of ProjectorStrings.
+
+        Returns:
+            A ProjectorSum.
+        """
+        if isinstance(terms, ProjectorString):
+            terms = [terms]
+        termdict: DefaultDict[FrozenSet[Tuple[raw_types.Qid, int]], value.Scalar] = defaultdict(
+            lambda: 0.0
+        )
+        for pstring in terms:
+            key = frozenset(pstring.projector_dict.items())
+            termdict[key] += pstring.coefficient
+        return cls(linear_dict=value.LinearDict(termdict))
+
+    def copy(self) -> 'ProjectorSum':
+        return ProjectorSum(self._linear_dict.copy())
+
+    def matrix(self, projector_qids: Optional[Iterable[raw_types.Qid]] = None) -> csr_matrix:
+        """Returns the matrix of self in computational basis of qubits.
+
+        Args:
+            projector_qids: Ordered collection of qubits that determine the subspace in which the
+                matrix representation of the ProjectorSum is to be computed. Qbits absent from
+                self.qubits are acted on by the identity. Defaults to the qubits of the
+                projector_dict.
+
+        Returns:
+            A sparse matrix that is the projection in the specified basis.
+        """
+        return sum(
+            coeff * _projector_string_from_projector_dict(vec).matrix(projector_qids)
+            for vec, coeff in self._linear_dict.items()
+        )
+
+    def expectation_from_state_vector(
+        self,
+        state_vector: np.ndarray,
+        qid_map: Mapping[raw_types.Qid, int],
+    ) -> float:
+        """Compute the expectation value of this ProjectorSum given a state vector.
+
+        Projects the state vector onto the sum of projectors and computes the expectation of the
+        measurements.
+
+        Args:
+            state_vector: An array representing a valid state vector.
+            qubit_map: A map from all qubits used in this ProjectorSum to the indices of the qubits
+                that `state_vector` is defined over.
+        Returns:
+            The expectation value of the input state.
+        """
+        return sum(
+            coeff
+            * _projector_string_from_projector_dict(vec).expectation_from_state_vector(
+                state_vector, qid_map
+            )
+            for vec, coeff in self._linear_dict.items()
+        )
+
+    def expectation_from_density_matrix(
+        self,
+        state: np.ndarray,
+        qid_map: Mapping[raw_types.Qid, int],
+    ) -> float:
+        """Expectation of the sum of projections from a density matrix.
+
+        Projects the density matrix onto the sum of projectors and computes the expectation of the
+        measurements.
+
+        Args:
+            state: An array representing a valid  density matrix.
+            qubit_map: A map from all qubits used in this ProjectorSum to the indices of the qubits
+                that `state_vector` is defined over.
+        Returns:
+            The expectation value of the input state.
+        """
+        return sum(
+            coeff
+            * _projector_string_from_projector_dict(vec).expectation_from_density_matrix(
+                state, qid_map
+            )
+            for vec, coeff in self._linear_dict.items()
+        )
+
+    def __iter__(self):
+        for vec, coeff in self._linear_dict.items():
+            yield _projector_string_from_projector_dict(vec, coeff)
+
+    def __len__(self) -> int:
+        return len(self._linear_dict)
+
+    def __truediv__(self, a: value.Scalar):
+        return self.__mul__(1 / a)
+
+    def __bool__(self) -> bool:
+        return bool(self._linear_dict)
+
+    def __iadd__(self, other: Union['ProjectorString', 'ProjectorSum']):
+        if isinstance(other, ProjectorString):
+            other = ProjectorSum.from_projector_strings(other)
+        elif not isinstance(other, ProjectorSum):
+            return NotImplemented
+        self._linear_dict += other._linear_dict
+        return self
+
+    def __add__(self, other: Union['ProjectorString', 'ProjectorSum']):
+        if isinstance(other, ProjectorString):
+            other = ProjectorSum.from_projector_strings(other)
+        elif not isinstance(other, ProjectorSum):
+            return NotImplemented
+        result = self.copy()
+        result += other
+        return result
+
+    def __isub__(self, other: Union['ProjectorString', 'ProjectorSum']):
+        if isinstance(other, ProjectorString):
+            other = ProjectorSum.from_projector_strings(other)
+        elif not isinstance(other, ProjectorSum):
+            return NotImplemented
+        self._linear_dict -= other._linear_dict
+        return self
+
+    def __sub__(self, other: Union['ProjectorString', 'ProjectorSum']):
+        if isinstance(other, ProjectorString):
+            other = ProjectorSum.from_projector_strings(other)
+        elif not isinstance(other, ProjectorSum):
+            return NotImplemented
+        result = self.copy()
+        result -= other
+        return result
+
+    def __neg__(self):
+        factory = type(self)
+        return factory(-self._linear_dict)
+
+    def __imul__(self, other: value.Scalar):
+        if not isinstance(other, numbers.Complex):
+            return NotImplemented
+        self._linear_dict *= other
+        return self
+
+    def __mul__(self, other: value.Scalar):
+        if not isinstance(other, numbers.Complex):
+            return NotImplemented
+        result = self.copy()
+        result *= other
+        return result
+
+    def __rmul__(self, other: value.Scalar):
+        if not isinstance(other, numbers.Complex):
+            return NotImplemented
+        result = self.copy()
+        result *= other
+        return result