Function to group qubit operators into tensor product basis sets (#494)

* Added method to split a QubitOperator into tensor product basis sets. Unit tests also added.

* Raises added to docstring of TPB set function.

* Minor formatting

* Refactored code with helper function. Added simple qubit operator tests

* Added more tests to for tensor product basis code using simple qubit operators

* Switch to using RandomState, remove MolecularData tests, switch order of additions clause.

* Removed randomize=False version of TPB code (now only random). Also removed associated tests.

* Removed unnecessary imports (MolecularData and OrderedDict)

* Switch from assertTrue to assertEqual

* Bug fix in test (one valid outcome was missing)

* Refactored if statement for additions to sub-operators.

* Removed unnecessary whitespace

* Added self to readme and notice

NOTICE

@@ -17,6 +17,7 @@ Pranav Gokhale (University of Chicago)
 Thomas Haener (ETH Zurich)
 Tarini Hardikar (Dartmouth)
 Vojtech Havlicek (Oxford)
+Oscar Higgott (University College London)
 Cupjin Huang (University of Michigan)
 Josh Izaac (Xanadu)
 Zhang Jiang (NASA)

README.rst

@@ -132,6 +132,7 @@ Authors
 `Thomas Häner <https://github.com/thomashaener>`__ (ETH Zurich),
 `Tarini Hardikar <https://github.com/TariniHardikar>`__ (Dartmouth),
 `Vojtĕch Havlíček <https://github.com/VojtaHavlicek>`__ (Oxford),
+`Oscar Higgott <https://github.com/oscarhiggott>`__ (University College London),
 `Cupjin Huang <https://github.com/pertoX4726>`__ (University of Michigan),
 `Josh Izaac <https://github.com/josh146>`__ (Xanadu),
 `Zhang Jiang <https://ti.arc.nasa.gov/profile/zjiang3>`__ (NASA),

src/openfermion/utils/__init__.py

@@ -32,7 +32,8 @@
                               is_hermitian, is_identity,
                               normal_ordered, prune_unused_indices,
                               reorder, up_then_down,
-                              load_operator, save_operator)
+                              load_operator, save_operator,
+                              group_into_tensor_product_basis_sets)
 
 from ._rdm_mapping_functions import (kronecker_delta,
                                      map_two_pdm_to_two_hole_dm,

src/openfermion/utils/_operator_utils.py

@@ -17,6 +17,7 @@
 import os
 import copy
 import marshal
+from numpy.random import RandomState
 
 import numpy
 from scipy.sparse import spmatrix
@@ -846,3 +847,79 @@ def normal_ordered(operator, hbar=1.):
         ordered_operator += order_fn(term, coefficient, **kwargs)
 
     return ordered_operator
+
+
+def _find_compatible_basis(term, bases):
+    for basis in bases:
+        basis_qubits = {op[0] for op in basis}
+        conflicts = ((i, P) for (i, P) in term
+                     if i in basis_qubits and (i, P) not in basis)
+        if any(conflicts):
+            continue
+        return basis
+    return None
+
+
+def group_into_tensor_product_basis_sets(operator, seed=None):
+    """
+    Split an operator (instance of QubitOperator) into `sub-operator`
+    QubitOperators, where each sub-operator has terms that are diagonal
+    in the same tensor product basis.
+
+    Each `sub-operator` can be measured using the same qubit post-rotations
+    in expectation estimation. Grouping into these tensor product basis
+    sets has been found to improve the efficiency of expectation estimation
+    significantly for some Hamiltonians in the context of
+    VQE (see section V(A) in the supplementary material of
+    https://arxiv.org/pdf/1704.05018v2.pdf). The more general problem
+    of grouping operators into commutitative groups is discussed in
+    section IV (B2) of https://arxiv.org/pdf/1509.04279v1.pdf. The
+    original input operator is the union of all output sub-operators,
+    and all sub-operators are disjoint (do not share any terms).
+
+    Args:
+        operator (QubitOperator): the operator that will be split into
+            sub-operators (tensor product basis sets).
+        seed (int): default None. Random seed used to initialize the
+            numpy.RandomState pseudo-random number generator.
+
+    Returns:
+        sub_operators (dict): a dictionary where each key defines a
+            tensor product basis, and each corresponding value is a
+            QubitOperator with terms that are all diagonal in
+            that basis.
+            **key** (tuple of tuples): Each key is a term, which defines
+                a tensor product basis. A term is a product of individual
+                factors; each factor is represented by a tuple of the form
+                (`index`, `action`), and these tuples are collected into a
+                larger tuple which represents the term as the product of
+                its factors. `action` is from the set {'X', 'Y', 'Z'} and
+                `index` is a non-negative integer corresponding to the
+                index of a qubit.
+            **value** (QubitOperator): A QubitOperator with terms that are
+                diagonal in the basis defined by the key it is stored in.
+
+    Raises:
+       TypeError: Operator of invalid type.
+    """
+    if not isinstance(operator, QubitOperator):
+        raise TypeError('Can only split QubitOperator into tensor product'
+                        ' basis sets. {} is not supported.'.format(
+                            type(operator).__name__))
+
+    sub_operators = {}
+    r = RandomState(seed)
+    for term, coefficient in operator.terms.items():
+        bases = list(sub_operators.keys())
+        r.shuffle(bases)
+        basis = _find_compatible_basis(term, bases)
+        if basis is None:
+            sub_operators[term] = QubitOperator(term, coefficient)
+        else:
+            sub_operator = sub_operators.pop(basis)
+            sub_operator += QubitOperator(term, coefficient)
+            additions = tuple(op for op in term if op not in basis)
+            basis = tuple(sorted(basis + additions, key=lambda factor: factor[0]))
+            sub_operators[basis] = sub_operator
+
+    return sub_operators

src/openfermion/utils/_operator_utils_test.py

@@ -843,3 +843,61 @@ def test_quad_triple(self):
     def test_exceptions(self):
         with self.assertRaises(TypeError):
             _ = normal_ordered(1)
+
+
+class GroupTensorProductBasisTest(unittest.TestCase):
+
+    def test_demo_qubit_operator(self):
+        for seed in [None, 0, 10000]:
+            op = QubitOperator('X0 Y1', 2.) + QubitOperator('X1 Y2', 3.j)
+            sub_operators = group_into_tensor_product_basis_sets(op, seed=seed)
+            expected = {((0, 'X'), (1, 'Y')): QubitOperator('X0 Y1', 2.),
+                        ((1, 'X'), (2, 'Y')): QubitOperator('X1 Y2', 3.j)}
+            self.assertEqual(sub_operators, expected)
+
+            op = QubitOperator('X0 Y1', 2.) + QubitOperator('Y1 Y2', 3.j)
+            sub_operators = group_into_tensor_product_basis_sets(op, seed=seed)
+            expected = {((0, 'X'), (1, 'Y'), (2, 'Y')): op}
+            self.assertEqual(sub_operators, expected)
+
+            op = QubitOperator('', 4.) + QubitOperator('X1', 2.j)
+            sub_operators = group_into_tensor_product_basis_sets(op, seed=seed)
+            expected = {((1, 'X'),): op}
+            self.assertEqual(sub_operators, expected)
+
+            op = (QubitOperator('X0 X1', 0.1) + QubitOperator('X1 X2', 2.j)
+                  + QubitOperator('Y2 Z3', 3.) + QubitOperator('X3 Z4', 5.))
+            sub_operators = group_into_tensor_product_basis_sets(op, seed=seed)
+            expected1 = {
+                ((0, 'X'), (1, 'X'), (2, 'X'),
+                 (3, 'X'), (4, 'Z')): (QubitOperator('X0 X1', 0.1)
+                                       + QubitOperator('X1 X2', 2.j)
+                                       + QubitOperator('X3 Z4', 5.)),
+                ((2, 'Y'), (3, 'Z')): QubitOperator('Y2 Z3', 3.)
+            }
+            expected2 = {
+                ((0, 'X'), (1, 'X'),
+                 (2, 'Y'), (3, 'Z')): (QubitOperator('X0 X1', 0.1)
+                                       + QubitOperator('Y2 Z3', 3.)),
+                ((1, 'X'), (2, 'X'),
+                 (3, 'X'), (4, 'Z')): (QubitOperator('X1 X2', 2.j)
+                                       + QubitOperator('X3 Z4', 5.))
+            }
+            self.assertTrue(sub_operators == expected1 or
+                            sub_operators == expected2)
+
+    def test_empty_qubit_operator(self):
+        sub_operators = group_into_tensor_product_basis_sets(QubitOperator())
+        self.assertTrue(sub_operators == {})
+
+    def test_fermion_operator_bad_type(self):
+        with self.assertRaises(TypeError):
+            _ = group_into_tensor_product_basis_sets(FermionOperator())
+
+    def test_boson_operator_bad_type(self):
+        with self.assertRaises(TypeError):
+            _ = group_into_tensor_product_basis_sets(BosonOperator())
+
+    def test_none_bad_type(self):
+        with self.assertRaises(TypeError):
+            _ = group_into_tensor_product_basis_sets(None)