Greedy term grouping (#754)

* Initial commit for greedy term grouping of Tomography Experiments

* Improved comments

* Updating function and variable names, adding to tests

* Using greedy term grouping by default

* Adding types to function inputs

* Removing unnecessary comments

* Fixing flake8 complaints

* Further fixing flake8 complaints

* Recalling my proper English grammar

* Listing allowed values for grouping method

* Making type of the input  explicit

* Improving doc-string for diagonal_basis_commutes(..)

* Changing variable names for consistency

* Making helper function (for greedy term grouping) private

* Blank lines between description and parameters in doc-string

* Giving the function 'diagonal_basis_commutes(..)' a better name

* Making the function 'get_diagonalizing_basis(..)' private

* Making _get_diagonalizing_basis(..) a bit more readable

* Fixing flake8 complaint

* Removing redundant code

* More readable code for _get_diagonalizing_basis(..)

* Adding private fuction extending the concept of being diagonal in the same tpb to ExperimentSettings

* Two 1q operators are diagonal in each others natural tpb if and only if they commute

* Refactoring to make grouping work with
ExperimentSettings more directly

* Minor flake8 fixes

* 'an' only comes before vowels

* _get_diagonalizing_basis makes _validate_all_diagonal_in_tpb do most of the work

* Updating doc-string for _all_qubits_diagonal_in_tpb

* Using already defined function to simplify code in other function

* Shortening variable names

* Preventing private function from mutating data structure

* Ensuring updated tpb is at least as large as the one associated with ExperimentSetting it's being updated for

* Re-organizing functions according to functionality

* Manually bringing in changes from master (merge conflicts are messy)

* Fixing failing test

* Remove unused import

* Don't import * from paulis (pulls in gunk since no __all__ defined)

* Move things back to where they once were

* pep8

* Split tests

pyquil/operator_estimation.py

@@ -5,15 +5,14 @@
 from json import JSONEncoder
 from math import pi
 from typing import List, Union, Iterable, Dict
-
 import networkx as nx
 import numpy as np
 from networkx.algorithms.approximation.clique import clique_removal
 
 from pyquil import Program
 from pyquil.api import QuantumComputer
 from pyquil.gates import *
-from pyquil.paulis import PauliTerm, is_identity
+from pyquil.paulis import PauliTerm, is_identity, sI
 
 if sys.version_info < (3, 7):
     from pyquil.external.dataclasses import dataclass
@@ -275,11 +274,13 @@ def _local_pauli_eig_meas(op, idx):
     raise ValueError(f'Unknown operation {op}')
 
 
-def _ops_diagonal_in_tpb(op_code1: str, op_code2: str):
+def _ops_commute(op_code1: str, op_code2: str):
     """
-    Given two op strings (I, X, Y, or Z) return whether they are diagonal in a tensor product basis.
+    Given two 1q op strings (I, X, Y, or Z), determine whether they commute
 
-    I.e. are they the same or is one of them 'I'.
+    :param op_code1: First operation
+    :param op_code2: Second operation
+    :return: Boolean specifying whether the two operations commute
     """
     if op_code1 not in ['X', 'Y', 'Z', 'I']:
         raise ValueError(f"Unknown op_code {op_code1}")
@@ -300,10 +301,53 @@ def _ops_diagonal_in_tpb(op_code1: str, op_code2: str):
 def _all_qubits_diagonal_in_tpb(op1: PauliTerm, op2: PauliTerm):
     """
     Compare all qubits between two PauliTerms to see if they are all diagonal in an
-    overall shared tensor product basis.
+    overall shared tensor product basis. More concretely, test if ``op1`` and
+    ``op2`` are diagonal in each others' "natural" tensor product basis.
+
+    Given some PauliTerm, the 'natural' tensor product basis (tpb) to
+    diagonalize this term is the one which diagonalizes each Pauli operator in the
+    product term-by-term.
+
+    For example, X(1) * Z(0) would be diagonal in the 'natural' tensor product basis
+    {(|0> +/- |1>)/Sqrt[2]} * {|0>, |1>}, whereas Z(1) * X(0) would be diagonal
+    in the 'natural' tpb {|0>, |1>} * {(|0> +/- |1>)/Sqrt[2]}. The two operators
+    commute but are not diagonal in each others 'natural' tpb (in fact, they are
+    anti-diagonal in each others 'natural' tpb). This function tests whether two
+    operators given as PauliTerms are both diagonal in each others 'natural' tpb.
+
+    Note that for the given example of X(1) * Z(0) and Z(1) * X(0), we can construct
+    the following basis which simultaneously diagonalizes both operators:
+
+      -- |0>' = |0> (|+>) + |1> (|->)
+      -- |1>' = |0> (|+>) - |1> (|->)
+      -- |2>' = |0> (|->) + |1> (|+>)
+      -- |3>' = |0> (-|->) + |1> (|+>)
+
+    In this basis, X Z looks like diag(1, -1, 1, -1), and Z X looks like diag(1, 1, -1, -1).
+    Notice however that this basis cannot be constructed with single-qubit operations, as each
+    of the basis vectors are entangled states.
+
+    :param op1: PauliTerm to check diagonality of in the natural tpb of ``op2``
+    :param op2: PauliTerm to check diagonality of in the natural tpb of ``op1``
+    :return: Boolean of diagonality in each others natural tpb
     """
-    all_qubits = set(op1.get_qubits()) | set(op2.get_qubits())
-    return all(_ops_diagonal_in_tpb(op1[q], op2[q]) for q in all_qubits)
+    all_qubits = set(op1.get_qubits()) & set(op2.get_qubits())
+    return all(_ops_commute(op1[q], op2[q]) for q in all_qubits)
+
+
+def _expt_settings_diagonal_in_tpb(es1: ExperimentSetting, es2: ExperimentSetting):
+    """
+    Extends the concept of being diagonal in the same tpb (see :py:func:_all_qubits_diagonal_in_tpb)
+    to ExperimentSettings, by determining if the pairs of in_operators and out_operators are
+    separately diagonal in the same tpb
+
+    :param es1: ExperimentSetting to check diagonality of in the natural tpb of ``es2``
+    :param es2: ExperimentSetting to check diagonality of in the natural tpb of ``es1``
+    :return: Boolean of diagonality in each others natural tpb
+    """
+    in_dtpb = _all_qubits_diagonal_in_tpb(es1.in_operator, es2.in_operator)
+    out_dtpb = _all_qubits_diagonal_in_tpb(es1.out_operator, es2.out_operator)
+    return in_dtpb and out_dtpb
 
 
 def construct_tpb_graph(experiments: TomographyExperiment):
@@ -327,20 +371,19 @@ def construct_tpb_graph(experiments: TomographyExperiment):
         if expt1 == expt2:
             continue
 
-        if (_all_qubits_diagonal_in_tpb(expt1.in_operator, expt2.in_operator)
-                and _all_qubits_diagonal_in_tpb(expt1.out_operator, expt2.out_operator)):
+        if _expt_settings_diagonal_in_tpb(expt1, expt2):
             g.add_edge(expt1, expt2)
 
     return g
 
 
-def group_experiments(experiments: TomographyExperiment) -> TomographyExperiment:
+def group_experiments_clique_removal(experiments: TomographyExperiment) -> TomographyExperiment:
     """
     Group experiments that are diagonal in a shared tensor product basis (TPB) to minimize number
-    of QPU runs.
+    of QPU runs, using a graph clique removal algorithm.
 
-    :param experiments: an tomography experiment
-    :return: an tomography experiment with all the same settings, just grouped according to shared
+    :param experiments: a tomography experiment
+    :return: a tomography experiment with all the same settings, just grouped according to shared
         TPBs.
     """
     g = construct_tpb_graph(experiments)
@@ -357,6 +400,122 @@ def group_experiments(experiments: TomographyExperiment) -> TomographyExperiment
     return TomographyExperiment(new_cliqs, program=experiments.program, qubits=experiments.qubits)
 
 
+def _validate_all_diagonal_in_tpb(ops: Iterable[PauliTerm]) -> Dict[int, str]:
+    """Each non-identity qubit should result in the same op_str among all operations. Return
+    said mapping.
+    """
+    mapping = dict()  # type: Dict[int, str]
+    for op in ops:
+        for idx, op_str in op:
+            if idx in mapping:
+                assert mapping[idx] == op_str, 'Improper grouping of operators'
+            else:
+                mapping[idx] = op_str
+    return mapping
+
+
+def _get_diagonalizing_basis(ops: Iterable[PauliTerm]) -> PauliTerm:
+    """
+    Find the Pauli Term with the most non-identity terms
+
+    :param ops: Iterable of PauliTerms to check
+    :return: The highest weight PauliTerm from the input iterable
+    """
+    # obtain qubit: operation mapping
+    dict_pt = _validate_all_diagonal_in_tpb(ops)
+    # convert this mapping to PauliTerm
+    pt = sI()
+    for q, op in dict_pt.items():
+        pt *= PauliTerm(op, q)
+    return pt
+
+
+def _max_tpb_overlap(tomo_expt: TomographyExperiment):
+    """
+    Given an input TomographyExperiment, provide a dictionary indicating which ExperimentSettings
+    share a tensor product basis
+
+    :param tomo_expt: TomographyExperiment, from which to group ExperimentSettings that share a tpb
+        and can be run together
+    :return: dictionary keyed with ExperimentSetting (specifying a tpb), and with each value being a
+            list of ExperimentSettings (diagonal in that tpb)
+    """
+    # initialize empty dictionary
+    diagonal_sets = {}
+    # loop through ExperimentSettings of the TomographyExperiment
+    for expt_setting in tomo_expt:
+        # no need to group already grouped TomographyExperiment
+        assert len(expt_setting) == 1, 'already grouped?'
+        expt_setting = expt_setting[0]
+        # calculate max overlap of expt_setting with keys of diagonal_sets
+        # keep track of whether a shared tpb was found
+        found_tpb = False
+        # loop through dict items
+        for es, es_list in diagonal_sets.items():
+            # update the dict value if es is diagonal in the same tpb as expt_setting
+            if _expt_settings_diagonal_in_tpb(es, expt_setting):
+                # shared tpb was found
+                found_tpb = True
+                # determine the updated list of ExperimentSettings
+                updated_es_list = es_list + [expt_setting]
+                # obtain the diagonalizing bases for both the updated in and out sets
+                diag_in_term = _get_diagonalizing_basis([expst.in_operator for expst in updated_es_list])
+                diag_out_term = _get_diagonalizing_basis([expst.out_operator for expst in updated_es_list])
+                assert len(diag_in_term) >= len(es.in_operator), \
+                    "Highest weight in-PauliTerm can't be smaller than the given in-PauliTerm"
+                assert len(diag_out_term) >= len(es.out_operator), \
+                    "Highest weight out-PauliTerm can't be smaller than the given out-PauliTerm"
+                # update the diagonalizing basis (key of dict) if necessary
+                if len(diag_in_term) > len(es.in_operator) or len(diag_out_term) > len(es.out_operator):
+                    del diagonal_sets[es]
+                    new_es = ExperimentSetting(diag_in_term, diag_out_term)
+                    diagonal_sets[new_es] = updated_es_list
+                else:
+                    diagonal_sets[es] = updated_es_list
+                break
+
+        if not found_tpb:
+            # made it through entire dict without finding any ExperimentSetting with shared tpb,
+            # so need to make a new item
+            diagonal_sets[expt_setting] = [expt_setting]
+
+    return diagonal_sets
+
+
+def group_experiments_greedy(tomo_expt: TomographyExperiment):
+    """
+    Greedy method to group ExperimentSettings in a given TomographyExperiment
+
+    :param tomo_expt: TomographyExperiment to group ExperimentSettings within
+    :return: TomographyExperiment, with grouped ExperimentSettings according to whether
+        it consists of PauliTerms diagonal in the same tensor product basis
+    """
+    diag_sets = _max_tpb_overlap(tomo_expt)
+    grouped_expt_settings_list = list(diag_sets.values())
+    grouped_tomo_expt = TomographyExperiment(grouped_expt_settings_list, program=tomo_expt.program,
+                                             qubits=tomo_expt.qubits)
+    return grouped_tomo_expt
+
+
+def group_experiments(experiments: TomographyExperiment, method: str = 'greedy') -> TomographyExperiment:
+    """
+    Group experiments that are diagonal in a shared tensor product basis (TPB) to minimize number
+    of QPU runs, using a specified method (greedy method by default)
+
+    :param experiments: a tomography experiment
+    :param method: method used for grouping; the allowed methods are one of
+        ['greedy', 'clique-removal']
+    :return: a tomography experiment with all the same settings, just grouped according to shared
+        TPBs.
+    """
+    allowed_methods = ['greedy', 'clique-removal']
+    assert method in allowed_methods, f"'method' should be one of {allowed_methods}."
+    if method == 'greedy':
+        return group_experiments_greedy(experiments)
+    elif method == 'clique-removal':
+        return group_experiments_clique_removal(experiments)
+
+
 @dataclass(frozen=True)
 class ExperimentResult:
     """An expectation and standard deviation for the measurement of one experiment setting
@@ -384,24 +543,10 @@ def serializable(self):
         }
 
 
-def _validate_all_diagonal_in_tpb(ops: Iterable[PauliTerm]) -> Dict[int, str]:
-    """Each non-identity qubit should result in the same op_str among all operations. Return
-    said mapping.
-    """
-    mapping = dict()  # type: Dict[int, str]
-    for op in ops:
-        for idx, op_str in op:
-            if idx in mapping:
-                assert mapping[idx] == op_str, 'Improper grouping of operators'
-            else:
-                mapping[idx] = op_str
-    return mapping
-
-
 def measure_observables(qc: QuantumComputer, tomo_experiment: TomographyExperiment, n_shots=1000,
                         progress_callback=None, active_reset=False):
     """
-    Measure all the observables in an TomographyExperiment.
+    Measure all the observables in a TomographyExperiment.
 
     :param qc: A QuantumComputer which can run quantum programs
     :param tomo_experiment: A suite of tomographic observables to measure

pyquil/tests/test_operator_estimation.py

@@ -10,8 +10,10 @@
 
 from pyquil.api import WavefunctionSimulator
 from pyquil.operator_estimation import ExperimentSetting, TomographyExperiment, to_json, read_json, \
-    _all_qubits_diagonal_in_tpb, group_experiments, ExperimentResult, measure_observables
-from pyquil.paulis import sI, sX, sY, sZ, PauliSum
+    _all_qubits_diagonal_in_tpb, group_experiments, ExperimentResult, measure_observables, \
+    _get_diagonalizing_basis, _max_tpb_overlap, group_experiments_greedy, \
+    _expt_settings_diagonal_in_tpb
+from pyquil.paulis import sI, sX, sY, sZ, PauliSum, PauliTerm
 from pyquil import Program, get_qc
 from pyquil.gates import *
 
@@ -121,6 +123,13 @@ def test_all_ops_belong_to_tpb():
             assert _all_qubits_diagonal_in_tpb(e1.in_operator, e2.in_operator)
             assert _all_qubits_diagonal_in_tpb(e1.out_operator, e2.out_operator)
 
+    assert _all_qubits_diagonal_in_tpb(sZ(0), sZ(0) * sZ(1))
+    assert _all_qubits_diagonal_in_tpb(sX(5), sZ(4))
+    assert not _all_qubits_diagonal_in_tpb(sX(0), sY(0) * sZ(2))
+    # this last example illustrates that a pair of commuting operators
+    # need not be diagonal in the same tpb
+    assert not _all_qubits_diagonal_in_tpb(sX(1) * sZ(0), sZ(1) * sX(0))
+
 
 def test_group_experiments():
     expts = [  # cf above, I removed the inner nesting. Still grouped visually
@@ -229,6 +238,76 @@ def test_no_complex_coeffs(forest):
         res = list(measure_observables(qc, suite))
 
 
+def test_get_diagonalizing_basis_1():
+    pauli_terms = [sZ(0), sX(1) * sZ(0), sY(2) * sX(1)]
+    assert _get_diagonalizing_basis(pauli_terms) == sY(2) * sX(1) * sZ(0)
+
+
+def test_get_diagonalizing_basis_2():
+    pauli_terms = [sZ(0), sX(1) * sZ(0), sY(2) * sX(1), sZ(5) * sI(3)]
+    assert _get_diagonalizing_basis(pauli_terms) == sZ(5) * sY(2) * sX(1) * sZ(0)
+
+
+def test_max_tpb_overlap_1():
+    tomo_expt_settings = [ExperimentSetting(sZ(1) * sX(0), sY(2) * sY(1)),
+                          ExperimentSetting(sX(2) * sZ(1), sY(2) * sZ(0))]
+    tomo_expt_program = Program(H(0), H(1), H(2))
+    tomo_expt_qubits = [0, 1, 2]
+    tomo_expt = TomographyExperiment(tomo_expt_settings, tomo_expt_program, tomo_expt_qubits)
+    expected_dict = {ExperimentSetting(sX(0) * sZ(1) * sX(2), sZ(0) * sY(1) * sY(2)):
+                         [ExperimentSetting(sZ(1) * sX(0), sY(2) * sY(1)),
+                          ExperimentSetting(sX(2) * sZ(1), sY(2) * sZ(0))]}
+    assert expected_dict == _max_tpb_overlap(tomo_expt)
+
+
+def test_max_tpb_overlap_2():
+    expt_setting = ExperimentSetting(PauliTerm.from_compact_str('(1+0j)*Z7Y8Z1Y4Z2Y5Y0X6'),
+                                     PauliTerm.from_compact_str('(1+0j)*Z4X8Y5X3Y7Y1'))
+    p = Program(H(0), H(1), H(2))
+    qubits = [0, 1, 2]
+    tomo_expt = TomographyExperiment([expt_setting], p, qubits)
+    expected_dict = {expt_setting: [expt_setting]}
+    assert expected_dict == _max_tpb_overlap(tomo_expt)
+
+
+def test_max_tpb_overlap_3():
+    # add another ExperimentSetting to the above
+    expt_setting = ExperimentSetting(PauliTerm.from_compact_str('(1+0j)*Z7Y8Z1Y4Z2Y5Y0X6'),
+                                     PauliTerm.from_compact_str('(1+0j)*Z4X8Y5X3Y7Y1'))
+    expt_setting2 = ExperimentSetting(sZ(7), sY(1))
+    p = Program(H(0), H(1), H(2))
+    qubits = [0, 1, 2]
+    tomo_expt2 = TomographyExperiment([expt_setting, expt_setting2], p, qubits)
+    expected_dict2 = {expt_setting: [expt_setting, expt_setting2]}
+    assert expected_dict2 == _max_tpb_overlap(tomo_expt2)
+
+
+def test_group_experiments_greedy():
+    ungrouped_tomo_expt = TomographyExperiment(
+        [[ExperimentSetting(PauliTerm.from_compact_str('(1+0j)*Z7Y8Z1Y4Z2Y5Y0X6'),
+                            PauliTerm.from_compact_str('(1+0j)*Z4X8Y5X3Y7Y1'))],
+         [ExperimentSetting(sZ(7), sY(1))]], program=Program(H(0), H(1), H(2)),
+        qubits=[0, 1, 2])
+    grouped_tomo_expt = group_experiments_greedy(ungrouped_tomo_expt)
+    expected_grouped_tomo_expt = TomographyExperiment(
+        [[ExperimentSetting(PauliTerm.from_compact_str('(1+0j)*Z7Y8Z1Y4Z2Y5Y0X6'),
+                            PauliTerm.from_compact_str('(1+0j)*Z4X8Y5X3Y7Y1')),
+          ExperimentSetting(sZ(7), sY(1))]],
+        program=Program(H(0), H(1), H(2)),
+        qubits=[0, 1, 2])
+    assert grouped_tomo_expt == expected_grouped_tomo_expt
+
+
+def test_expt_settings_diagonal_in_tpb():
+    expt_setting1 = ExperimentSetting(sZ(1) * sX(0), sY(1) * sZ(0))
+    expt_setting2 = ExperimentSetting(sY(2) * sZ(1), sZ(2) * sY(1))
+    assert _expt_settings_diagonal_in_tpb(expt_setting1, expt_setting2)
+    expt_setting3 = ExperimentSetting(sX(2) * sZ(1), sZ(2) * sY(1))
+    expt_setting4 = ExperimentSetting(sY(2) * sZ(1), sX(2) * sY(1))
+    assert not _expt_settings_diagonal_in_tpb(expt_setting2, expt_setting3)
+    assert not _expt_settings_diagonal_in_tpb(expt_setting2, expt_setting4)
+
+
 def test_identity(forest):
     qc = get_qc('2q-qvm')
     suite = TomographyExperiment([ExperimentSetting(sI(), 0.123 * sI(0))],