Add new passmanager.replace(index, ...) method (#3004)

* replace method

* error handle

* lint

* changelog

* lint

* release note

* key->index

* error message

* lint

* release note

* iamges

* replace images for text

qiskit/transpiler/passmanager.py

@@ -127,25 +127,68 @@ def append(self, passes, max_iteration=None, **flow_controller_conditions):
         Raises:
             TranspilerError: if a pass in passes is not a proper pass.
         """
-        passset_options = {'max_iteration': max_iteration}
+        options = self._join_options({'max_iteration': max_iteration})
 
-        options = self._join_options(passset_options)
+        passes = PassManager._normalize_passes(passes)
 
+        flow_controller_conditions = self._normalize_flow_controller(flow_controller_conditions)
+
+        self.working_list.append(
+            FlowController.controller_factory(passes, options, **flow_controller_conditions))
+
+    def replace(self, index, passes, max_iteration=None, **flow_controller_conditions):
+        """Replace a particular pass in the scheduler
+
+        Args:
+            index (int): Pass index to replace, based on the position in passes().
+            passes (list[BasePass] or BasePass): pass(es) to be added to schedule
+            max_iteration (int): max number of iterations of passes. Default: 1000
+            flow_controller_conditions (kwargs): See add_flow_controller(): Dictionary of
+            control flow plugins. Default:
+
+                * do_while (callable property_set -> boolean): The passes repeat until the
+                  callable returns False.
+                  Default: `lambda x: False # i.e. passes run once`
+
+                * condition (callable property_set -> boolean): The passes run only if the
+                  callable returns True.
+                  Default: `lambda x: True # i.e. passes run`
+        Raises:
+            TranspilerError: if a pass in passes is not a proper pass.
+        """
+        options = self._join_options({'max_iteration': max_iteration})
+
+        passes = PassManager._normalize_passes(passes)
+
+        flow_controller_conditions = self._normalize_flow_controller(flow_controller_conditions)
+
+        controller = FlowController.controller_factory(passes, options,
+                                                       **flow_controller_conditions)
+        try:
+            self.working_list[index] = controller
+        except IndexError:
+            raise TranspilerError('Index to replace %s does not exists' % index)
+
+    def __setitem__(self, index, item):
+        self.replace(index, item)
+
+    @staticmethod
+    def _normalize_passes(passes):
         if isinstance(passes, BasePass):
             passes = [passes]
 
         for pass_ in passes:
             if not isinstance(pass_, BasePass):
                 raise TranspilerError('%s is not a pass instance' % pass_.__class__)
+        return passes
 
-        for name, param in flow_controller_conditions.items():
+    def _normalize_flow_controller(self, flow_controller):
+        for name, param in flow_controller.items():
             if callable(param):
-                flow_controller_conditions[name] = partial(param, self.fenced_property_set)
+                flow_controller[name] = partial(param, self.fenced_property_set)
             else:
                 raise TranspilerError('The flow controller parameter %s is not callable' % name)
-
-        self.working_list.append(
-            FlowController.controller_factory(passes, options, **flow_controller_conditions))
+        return flow_controller
 
     def reset(self):
         """Reset the pass manager instance"""

qiskit/visualization/pass_manager_visualization.py

@@ -102,10 +102,10 @@ def pass_manager_drawer(pass_manager, filename, style=None, raw=False):
 
     prev_node = None
 
-    for controller_group in passes:
+    for index, controller_group in enumerate(passes):
 
         # label is the name of the flow controller (without the word controller)
-        label = controller_group['type'].__name__.replace('Controller', '')
+        label = "[%s] %s" % (index, controller_group['type'].__name__.replace('Controller', ''))
 
         # create the subgraph for this controller
         subgraph = pydot.Cluster(str(component_id), label=label, fontname='helvetica')

releasenotes/notes/passmanager_replace-d89e2cc46517d917.yaml

@@ -0,0 +1,47 @@
+---
+features:
+  - |
+    A given pass manager now can be edited with the new method `replace`. This method allows to
+    replace a particular stage in a pass manager, which can be handy when dealing with preset
+    pass managers. For example, let's edit the layout selector of the pass manager used at
+    optimization level 0:
+
+    .. code-block:: python
+      from qiskit.transpiler.preset_passmanagers.level0 import level_0_pass_manager
+      from qiskit.transpiler.transpile_config import TranspileConfig
+
+      pass_manager = level_0_pass_manager(TranspileConfig(coupling_map=CouplingMap([[0,1]])))
+
+      pass_manager.draw()
+
+    .. code-block::
+      [0] FlowLinear: SetLayout
+      [1] Conditional: TrivialLayout
+      [2] FlowLinear: FullAncillaAllocation, EnlargeWithAncilla, ApplyLayout
+      [3] FlowLinear: Unroller
+
+    The layout selection is set in the stage `[1]`. Let's replace it with `DenseLayout`:
+
+    .. code-block:: python
+      from qiskit.transpiler.passes import DenseLayout
+
+      pass_manager.replace(1, DenseLayout(coupling_map), condition=lambda property_set: not property_set['layout'])
+      pass_manager.draw()
+
+    .. code-block::
+      [0] FlowLinear: SetLayout
+      [1] Conditional: DenseLayout
+      [2] FlowLinear: FullAncillaAllocation, EnlargeWithAncilla, ApplyLayout
+      [3] FlowLinear: Unroller
+
+    If you want to replace it without any condition, you can use set-item shortcut:
+
+    .. code-block:: python
+      pass_manager[1] = DenseLayout(coupling_map)
+      pass_manager.draw()
+
+    .. code-block::
+      [0] FlowLinear: SetLayout
+      [1] FlowLinear: DenseLayout
+      [2] FlowLinear: FullAncillaAllocation, EnlargeWithAncilla, ApplyLayout
+      [3] FlowLinear: Unroller

test/python/transpiler/test_pass_scheduler.py

@@ -556,6 +556,7 @@ def test_conditional_and_loop(self):
 
 class StreamHandlerRaiseException(StreamHandler):
     """Handler class that will raise an exception on formatting errors."""
+
     def handleError(self, record):
         raise sys.exc_info()
 
@@ -720,5 +721,65 @@ def test_fixed_point_twice(self):
         self.assertScheduler(self.circuit, self.passmanager, expected)
 
 
+class TestPassManagerReplace(SchedulerTestCase):
+    """Test PassManager.replace"""
+
+    def setUp(self):
+        self.passmanager = PassManager()
+        self.circuit = QuantumCircuit(QuantumRegister(1))
+
+    def test_replace0(self):
+        """ Test passmanager.replace(0, ...)."""
+        self.passmanager.append(PassC_TP_RA_PA())  # Request: PassA / Preserves: PassA
+        self.passmanager.append(PassB_TP_RA_PA())  # Request: PassA / Preserves: PassA
+
+        self.passmanager.replace(0, PassB_TP_RA_PA())
+
+        expected = ['run transformation pass PassA_TP_NR_NP',
+                    'run transformation pass PassB_TP_RA_PA']
+        self.assertScheduler(self.circuit, self.passmanager, expected)
+
+    def test_replace1(self):
+        """ Test passmanager.replace(1, ...)."""
+        self.passmanager.append(PassC_TP_RA_PA())  # Request: PassA / Preserves: PassA
+        self.passmanager.append(PassB_TP_RA_PA())  # Request: PassA / Preserves: PassA
+
+        self.passmanager.replace(1, PassC_TP_RA_PA())
+
+        expected = ['run transformation pass PassA_TP_NR_NP',
+                    'run transformation pass PassC_TP_RA_PA']
+        self.assertScheduler(self.circuit, self.passmanager, expected)
+
+    def test_setitem(self):
+        """ Test passmanager[1] = ..."""
+        self.passmanager.append(PassC_TP_RA_PA())  # Request: PassA / Preserves: PassA
+        self.passmanager.append(PassB_TP_RA_PA())  # Request: PassA / Preserves: PassA
+
+        self.passmanager[1] = PassC_TP_RA_PA()
+
+        expected = ['run transformation pass PassA_TP_NR_NP',
+                    'run transformation pass PassC_TP_RA_PA']
+        self.assertScheduler(self.circuit, self.passmanager, expected)
+
+    def test_replace_with_conditional(self):
+        """ Replace a pass with a conditional pass. """
+        self.passmanager.append(PassE_AP_NR_NP(False))
+        self.passmanager.append(PassB_TP_RA_PA())
+
+        self.passmanager.replace(1, PassA_TP_NR_NP(),
+                                 condition=lambda property_set: property_set['property'])
+
+        expected = ['run analysis pass PassE_AP_NR_NP',
+                    'set property as False']
+        self.assertScheduler(self.circuit, self.passmanager, expected)
+
+    def test_replace_error(self):
+        """ Replace a non-existing index. """
+        self.passmanager.append(PassB_TP_RA_PA())
+
+        with self.assertRaises(TranspilerError):
+            self.passmanager.replace(99, PassA_TP_NR_NP())
+
+
 if __name__ == '__main__':
     unittest.main()