Refactor drop_negligible into an OptimizationPass using batch_remove (#695)

Author: Strilanc

cirq/circuits/drop_negligible.py

@@ -14,13 +14,17 @@
 
 """An optimization pass that removes operations with tiny effects."""
 
-from typing import Optional
+from typing import TYPE_CHECKING
 
 from cirq import ops, extension
 from cirq.circuits import optimization_pass, circuit as _circuit
 
+if TYPE_CHECKING:
+    # pylint: disable=unused-import
+    from typing import List, Tuple
 
-class DropNegligible(optimization_pass.PointOptimizer):
+
+class DropNegligible(optimization_pass.OptimizationPass):
     """An optimization pass that removes operations with tiny effects."""
 
     def __init__(self,
@@ -29,18 +33,12 @@ def __init__(self,
         self.tolerance = tolerance
         self.extensions = extensions or extension.Extensions()
 
-    def optimization_at(
-            self,
-            circuit: _circuit.Circuit,
-            index: int,
-            op: ops.Operation
-            ) -> Optional[optimization_pass.PointOptimizationSummary]:
-
-        bounded = self.extensions.try_cast(ops.BoundedEffect, op)
-        if bounded is None or bounded.trace_distance_bound() > self.tolerance:
-            return None
-
-        return optimization_pass.PointOptimizationSummary(
-            clear_span=1,
-            new_operations=(),
-            clear_qubits=op.qubits)
+    def optimize_circuit(self, circuit: _circuit.Circuit) -> None:
+        deletions = []  # type: List[Tuple[int, ops.Operation]]
+        for moment_index, moment in enumerate(circuit):
+            for op in moment.operations:
+                bounded = self.extensions.try_cast(ops.BoundedEffect, op)
+                if (bounded is not None and
+                        bounded.trace_distance_bound() <= self.tolerance):
+                    deletions.append((moment_index, op))
+        circuit.batch_remove(deletions)

cirq/circuits/drop_negligible_test.py

@@ -27,7 +27,6 @@ def test_leaves_big():
     drop = cirq.DropNegligible(0.001)
     a = cirq.NamedQubit('a')
     circuit = cirq.Circuit([cirq.Moment([cirq.Z(a)**0.1])])
-    assert drop.optimization_at(circuit, 0, circuit.operation_at(a, 0)) is None
 
     assert_optimizes(optimizer=drop,
                      initial_circuit=circuit,
@@ -39,11 +38,6 @@ def test_clears_small():
     a = cirq.NamedQubit('a')
     circuit = cirq.Circuit([cirq.Moment([cirq.Z(a)**0.000001])])
 
-    assert (drop.optimization_at(circuit, 0, circuit.operation_at(a, 0)) ==
-            cirq.PointOptimizationSummary(clear_span=1,
-                                          clear_qubits=[a],
-                                          new_operations=[]))
-
     assert_optimizes(optimizer=drop,
                      initial_circuit=circuit,
                      expected_circuit=cirq.Circuit([cirq.Moment()]))