Fix unexpected behavior of symmetry_conserving_bravyi_kitaev (quantum…

…lib#734)

* Fix unexpected behavior of symmetry_conserving_bravyi_kitaev

* Update docstring

* fix lint errors

Co-authored-by: Nicholas Rubin <[email protected]>

src/openfermion/transforms/opconversions/remove_symmetry_qubits.py

@@ -14,9 +14,10 @@
     arXiv:1701.08213 and Phys. Rev. X 6, 031007.
 """
 
+import copy
+
 from openfermion.ops.operators import FermionOperator
 from openfermion.transforms.opconversions import bravyi_kitaev_tree, reorder
-from openfermion.transforms.repconversions import prune_unused_indices
 from openfermion.utils.indexing import up_then_down
 
 
@@ -61,12 +62,12 @@ def symmetry_conserving_bravyi_kitaev(fermion_hamiltonian, active_orbitals,
     """
     # Catch errors if inputs are of wrong type.
     if type(fermion_hamiltonian) is not FermionOperator:
-        raise ValueError('Supplied operator should be an instance '
-                         'of FermionOperator class')
+        raise ValueError(
+            "Supplied operator should be an instance of FermionOperator class")
     if type(active_orbitals) is not int:
-        raise ValueError('Number of active orbitals should be an integer.')
+        raise ValueError("Number of active orbitals should be an integer.")
     if type(active_fermions) is not int:
-        raise ValueError('Number of active fermions should be an integer.')
+        raise ValueError("Number of active fermions should be an integer.")
 
     # Arrange spins up then down, then BK map to qubit Hamiltonian.
     fermion_hamiltonian_reorder = reorder(fermion_hamiltonian,
@@ -98,7 +99,8 @@ def symmetry_conserving_bravyi_kitaev(fermion_hamiltonian, active_orbitals,
     qubit_hamiltonian = edit_hamiltonian_for_spin(qubit_hamiltonian,
                                                   active_orbitals / 2,
                                                   parity_middle_orb)
-    qubit_hamiltonian = prune_unused_indices(qubit_hamiltonian)
+    qubit_hamiltonian = remove_indices(qubit_hamiltonian,
+                                       (active_orbitals / 2, active_orbitals))
 
     return qubit_hamiltonian
 
@@ -110,9 +112,9 @@ def edit_hamiltonian_for_spin(qubit_hamiltonian, spin_orbital, orbital_parity):
     for term, coefficient in qubit_hamiltonian.terms.items():
         # If Z acts on the specified orbital, precompute its effect and
         # remove it from the Hamiltonian.
-        if (spin_orbital - 1, 'Z') in term:
+        if (spin_orbital - 1, "Z") in term:
             new_coefficient = coefficient * orbital_parity
-            new_term = tuple(i for i in term if i != (spin_orbital - 1, 'Z'))
+            new_term = tuple(i for i in term if i != (spin_orbital - 1, "Z"))
             # Make sure to combine terms comprised of the same operators.
             if new_qubit_dict.get(new_term) is None:
                 new_qubit_dict[new_term] = new_coefficient
@@ -131,3 +133,32 @@ def edit_hamiltonian_for_spin(qubit_hamiltonian, spin_orbital, orbital_parity):
     qubit_hamiltonian.compress()
 
     return qubit_hamiltonian
+
+
+def remove_indices(symbolic_operator, indices):
+    """Returns the symbolic operator from which the operator with the specified
+    index was removed.
+
+    Args:
+        symbolic_operator: An instance of the SymbolicOperator class.
+        indices: A sequence of Int type object. The indices to be removed.
+
+    Returns:
+        The symbolic operator. The removed indices will be filled by shifting.
+    """
+    map = {}
+
+    def new_index(index):
+        if index in map:
+            return map[index]
+        map[index] = index - len(list(i for i in indices if (i - 1) < index))
+        return map[index]
+
+    new_operator = copy.deepcopy(symbolic_operator)
+    new_operator.terms.clear()
+
+    for term in symbolic_operator.terms:
+        new_term = [(new_index(op[0]), op[1]) for op in term]
+        new_operator.terms[tuple(new_term)] = symbolic_operator.terms[term]
+
+    return new_operator