Allow full gates to be passed into get method of InstructionScheduleMap

qiskit/pulse/instruction_schedule_map.py

@@ -33,6 +33,7 @@
 from qiskit.circuit.parameterexpression import ParameterExpression
 from qiskit.pulse.exceptions import PulseError
 from qiskit.pulse.schedule import ParameterizedSchedule, Schedule
+from qiskit.circuit.instruction import Instruction
 
 
 class InstructionScheduleMap():
@@ -68,7 +69,9 @@ def instructions(self) -> List[str]:
         """
         return list(self._map.keys())
 
-    def qubits_with_instruction(self, instruction: str) -> List[Union[int, Tuple[int]]]:
+    def qubits_with_instruction(self,
+                                instruction: Union[str, Instruction]) -> List[Union[int,
+                                                                                    Tuple[int]]]:
         """Return a list of the qubits for which the given instruction is defined. Single qubit
         instructions return a flat list, and multiqubit instructions return a list of ordered
         tuples.
@@ -83,6 +86,7 @@ def qubits_with_instruction(self, instruction: str) -> List[Union[int, Tuple[int
         Raises:
             PulseError: If the instruction is not found.
         """
+        instruction = _get_instruction_string(instruction)
         if instruction not in self._map:
             return []
         return [qubits[0] if len(qubits) == 1 else qubits
@@ -105,7 +109,7 @@ def qubit_instructions(self, qubits: Union[int, Iterable[int]]) -> List[str]:
             return list(self._qubit_instructions[_to_tuple(qubits)])
         return []
 
-    def has(self, instruction: str, qubits: Union[int, Iterable[int]]) -> bool:
+    def has(self, instruction: Union[str, Instruction], qubits: Union[int, Iterable[int]]) -> bool:
         """Is the instruction defined for the given qubits?
 
         Args:
@@ -115,10 +119,13 @@ def has(self, instruction: str, qubits: Union[int, Iterable[int]]) -> bool:
         Returns:
             True iff the instruction is defined.
         """
+        instruction = _get_instruction_string(instruction)
         return instruction in self._map and \
             _to_tuple(qubits) in self._map[instruction]
 
-    def assert_has(self, instruction: str, qubits: Union[int, Iterable[int]]) -> None:
+    def assert_has(self,
+                   instruction: Union[str, Instruction],
+                   qubits: Union[int, Iterable[int]]) -> None:
         """Error if the given instruction is not defined.
 
         Args:
@@ -128,6 +135,7 @@ def assert_has(self, instruction: str, qubits: Union[int, Iterable[int]]) -> Non
         Raises:
             PulseError: If the instruction is not defined on the qubits.
         """
+        instruction = _get_instruction_string(instruction)
         if not self.has(instruction, _to_tuple(qubits)):
             if instruction in self._map:
                 raise PulseError("Operation '{inst}' exists, but is only defined for qubits "
@@ -138,22 +146,23 @@ def assert_has(self, instruction: str, qubits: Union[int, Iterable[int]]) -> Non
                              "system.".format(inst=instruction))
 
     def get(self,
-            instruction: str,
+            instruction: Union[str, Instruction],
             qubits: Union[int, Iterable[int]],
             *params: Union[int, float, complex, ParameterExpression],
             **kwparams: Union[int, float, complex, ParameterExpression]) -> Schedule:
         """Return the defined :py:class:`~qiskit.pulse.Schedule` for the given instruction on
         the given qubits.
 
         Args:
-            instruction: Name of the instruction.
+            instruction: Name of the instruction or the instruction itself.
             qubits: The qubits for the instruction.
             *params: Command parameters for generating the output schedule.
             **kwparams: Keyworded command parameters for generating the schedule.
 
         Returns:
             The Schedule defined for the input.
         """
+        instruction = _get_instruction_string(instruction)
         self.assert_has(instruction, qubits)
         schedule_generator = self._map[instruction].get(_to_tuple(qubits))
 
@@ -163,7 +172,7 @@ def get(self,
         return schedule_generator
 
     def add(self,
-            instruction: str,
+            instruction: Union[str, Instruction],
             qubits: Union[int, Iterable[int]],
             schedule: Union[Schedule, Callable[..., Schedule]]) -> None:
         """Add a new known instruction for the given qubits and its mapping to a pulse schedule.
@@ -176,6 +185,8 @@ def add(self,
         Raises:
             PulseError: If the qubits are provided as an empty iterable.
         """
+        instruction = _get_instruction_string(instruction)
+
         qubits = _to_tuple(qubits)
         if qubits == ():
             raise PulseError("Cannot add definition {} with no target qubits.".format(instruction))
@@ -185,13 +196,16 @@ def add(self,
         self._map[instruction][qubits] = schedule
         self._qubit_instructions[qubits].add(instruction)
 
-    def remove(self, instruction: str, qubits: Union[int, Iterable[int]]) -> None:
+    def remove(self,
+               instruction: Union[str, Instruction],
+               qubits: Union[int, Iterable[int]]) -> None:
         """Remove the given instruction from the listing of instructions defined in self.
 
         Args:
             instruction: The name of the instruction to add.
             qubits: The qubits which the instruction applies to.
         """
+        instruction = _get_instruction_string(instruction)
         qubits = _to_tuple(qubits)
         self.assert_has(instruction, qubits)
         self._map[instruction].pop(qubits)
@@ -202,7 +216,7 @@ def remove(self, instruction: str, qubits: Union[int, Iterable[int]]) -> None:
             self._qubit_instructions.pop(qubits)
 
     def pop(self,
-            instruction: str,
+            instruction: Union[str, Instruction],
             qubits: Union[int, Iterable[int]],
             *params: Union[int, float, complex, ParameterExpression],
             **kwparams: Union[int, float, complex, ParameterExpression]) -> Schedule:
@@ -218,11 +232,14 @@ def pop(self,
         Returns:
             The Schedule defined for the input.
         """
+        instruction = _get_instruction_string(instruction)
         schedule = self.get(instruction, qubits, *params, **kwparams)
         self.remove(instruction, qubits)
         return schedule
 
-    def get_parameters(self, instruction: str, qubits: Union[int, Iterable[int]]) -> Tuple[str]:
+    def get_parameters(self,
+                       instruction: Union[str, Instruction],
+                       qubits: Union[int, Iterable[int]]) -> Tuple[str]:
         """Return the list of parameters taken by the given instruction on the given qubits.
 
         Args:
@@ -232,6 +249,8 @@ def get_parameters(self, instruction: str, qubits: Union[int, Iterable[int]]) ->
         Returns:
             The names of the parameters required by the instruction.
         """
+        instruction = _get_instruction_string(instruction)
+
         self.assert_has(instruction, qubits)
         schedule_generator = self._map[instruction][_to_tuple(qubits)]
         if isinstance(schedule_generator, ParameterizedSchedule):
@@ -267,3 +286,14 @@ def _to_tuple(values: Union[int, Iterable[int]]) -> Tuple[int, ...]:
         return tuple(values)
     except TypeError:
         return (values,)
+
+
+def _get_instruction_string(inst: Union[str, Instruction]):
+    if isinstance(inst, str):
+        return inst
+    else:
+        try:
+            return inst.name
+        except AttributeError:
+            raise PulseError('Input "inst" has no attribute "name".'
+                             'This should be a circuit "Instruction".')

qiskit/scheduler/lowering.py

@@ -132,7 +132,7 @@ def get_measure_schedule(qubit_mem_slots: Dict[int, int]) -> CircuitPulseDef:
 
             try:
                 circ_pulse_defs.append(
-                    CircuitPulseDef(schedule=inst_map.get(inst.name, inst_qubits, *inst.params),
+                    CircuitPulseDef(schedule=inst_map.get(inst, inst_qubits, *inst.params),
                                     qubits=inst_qubits))
             except PulseError:
                 raise QiskitError("Operation '{}' on qubit(s) {} not supported by the backend "

test/python/pulse/test_instruction_schedule_map.py

@@ -16,6 +16,7 @@
 import numpy as np
 
 import qiskit.pulse.library as library
+from qiskit.circuit.library.standard_gates import U1Gate, U3Gate, CXGate, XGate
 from qiskit.circuit.parameter import Parameter
 from qiskit.circuit.parameterexpression import ParameterExpression
 from qiskit.pulse import (InstructionScheduleMap, Play, PulseError, Schedule,
@@ -121,10 +122,9 @@ def test_get(self):
         sched = Schedule()
         sched.append(Play(Waveform(np.ones(5)), DriveChannel(0)))
         inst_map = InstructionScheduleMap()
+        inst_map.add('x', 0, sched)
 
-        inst_map.add('u1', 0, sched)
-
-        self.assertEqual(sched, inst_map.get('u1', (0,)))
+        self.assertEqual(sched, inst_map.get('x', (0,)))
 
     def test_remove(self):
         """Test removing a defined operation and removing an undefined operation."""
@@ -152,6 +152,126 @@ def test_pop(self):
         with self.assertRaises(PulseError):
             inst_map.pop('not_there', (0,))
 
+    def test_add_gate(self):
+        """Test add, and that errors are raised when expected."""
+        sched = Schedule()
+        sched.append(Play(Waveform(np.ones(5)), DriveChannel(0)))
+        inst_map = InstructionScheduleMap()
+
+        inst_map.add(U1Gate(0), 1, sched)
+        inst_map.add(U1Gate(0), 0, sched)
+
+        self.assertIn('u1', inst_map.instructions)
+        self.assertEqual(inst_map.qubits_with_instruction(U1Gate(0)), [0, 1])
+        self.assertTrue('u1' in inst_map.qubit_instructions(0))
+
+        with self.assertRaises(PulseError):
+            inst_map.add(U1Gate(0), (), sched)
+        with self.assertRaises(PulseError):
+            inst_map.add(U1Gate(0), 1, "not a schedule")
+
+    def test_instructions_gate(self):
+        """Test `instructions`."""
+        sched = Schedule()
+        inst_map = InstructionScheduleMap()
+
+        inst_map.add(U1Gate(0), 1, sched)
+        inst_map.add(U3Gate(0, 0, 0), 0, sched)
+
+        instructions = inst_map.instructions
+        for inst in ['u1', 'u3']:
+            self.assertTrue(inst in instructions)
+
+    def test_has_gate(self):
+        """Test `has` and `assert_has`."""
+        sched = Schedule()
+        inst_map = InstructionScheduleMap()
+
+        inst_map.add(U1Gate(0), (0,), sched)
+        inst_map.add(CXGate(), [0, 1], sched)
+
+        self.assertTrue(inst_map.has(U1Gate(0), [0]))
+        self.assertTrue(inst_map.has(CXGate(), (0, 1)))
+        with self.assertRaises(PulseError):
+            inst_map.assert_has('dne', [0])
+        with self.assertRaises(PulseError):
+            inst_map.assert_has(CXGate(), 100)
+
+    def test_has_from_mock_gate(self):
+        """Test `has` and `assert_has` from mock data."""
+        inst_map = FakeOpenPulse2Q().defaults().instruction_schedule_map
+        self.assertTrue(inst_map.has(U1Gate(0), [0]))
+        self.assertTrue(inst_map.has(CXGate(), (0, 1)))
+        self.assertTrue(inst_map.has(U3Gate(0, 0, 0), 0))
+        self.assertTrue(inst_map.has('measure', [0, 1]))
+        self.assertFalse(inst_map.has(U1Gate(0), [0, 1]))
+        with self.assertRaises(PulseError):
+            inst_map.assert_has('dne', [0])
+        with self.assertRaises(PulseError):
+            inst_map.assert_has(CXGate(), 100)
+
+    def test_qubits_with_instruction_gate(self):
+        """Test `qubits_with_instruction`."""
+        sched = Schedule()
+        inst_map = InstructionScheduleMap()
+
+        inst_map.add(U1Gate(0), (0,), sched)
+        inst_map.add(U1Gate(0), (1,), sched)
+        inst_map.add(CXGate(), [0, 1], sched)
+
+        self.assertEqual(inst_map.qubits_with_instruction(U1Gate(0)), [0, 1])
+        self.assertEqual(inst_map.qubits_with_instruction(CXGate()), [(0, 1)])
+        self.assertEqual(inst_map.qubits_with_instruction('none'), [])
+
+    def test_qubit_instructions_gate(self):
+        """Test `qubit_instructions`."""
+        sched = Schedule()
+        inst_map = InstructionScheduleMap()
+
+        inst_map.add(U1Gate(0), (0,), sched)
+        inst_map.add(U1Gate(0), (1,), sched)
+        inst_map.add(CXGate(), [0, 1], sched)
+
+        self.assertEqual(inst_map.qubit_instructions(0), ['u1'])
+        self.assertEqual(inst_map.qubit_instructions(1), ['u1'])
+        self.assertEqual(inst_map.qubit_instructions((0, 1)), ['cx'])
+        self.assertEqual(inst_map.qubit_instructions(10), [])
+
+    def test_get_gate(self):
+        """Test `get`."""
+        sched = Schedule()
+        sched.append(Play(Waveform(np.ones(5)), DriveChannel(0)))
+        inst_map = InstructionScheduleMap()
+        inst_map.add(XGate(), 0, sched)
+
+        self.assertEqual(sched, inst_map.get(XGate(), (0,)))
+
+    def test_remove_gate(self):
+        """Test removing a defined operation and removing an undefined operation."""
+        sched = Schedule()
+        inst_map = InstructionScheduleMap()
+
+        inst_map.add('tmp', 0, sched)
+        inst_map.remove('tmp', 0)
+        self.assertFalse(inst_map.has('tmp', 0))
+        with self.assertRaises(PulseError):
+            inst_map.remove('not_there', (0,))
+        self.assertFalse('tmp' in inst_map.qubit_instructions(0))
+
+    def test_pop_gate(self):
+        """Test pop with default."""
+        sched = Schedule()
+        inst_map = InstructionScheduleMap()
+
+        inst_map.add(XGate(), 100, sched)
+        self.assertEqual(inst_map.pop(XGate(), 100), sched)
+        self.assertFalse(inst_map.has(XGate(), 100))
+
+        self.assertEqual(inst_map.qubit_instructions(100), [])
+        self.assertEqual(inst_map.qubits_with_instruction(XGate()), [])
+        with self.assertRaises(PulseError):
+            inst_map.pop('not_there', (0,))
+
     def test_sequenced_parameterized_schedule(self):
         """Test parametrized schedule consists of multiple instruction. """