Add results, state vector, amplitude, probability

CODEOWNERS

@@ -5,3 +5,4 @@
 # will be requested for review when someone opens a pull request.
 *   @floralph
 *   @speller26
+*   @avawang1

src/braket/circuits/__init__.py

@@ -26,3 +26,5 @@
 from braket.circuits.operator import Operator  # noqa: F401
 from braket.circuits.qubit import Qubit, QubitInput  # noqa: F401
 from braket.circuits.qubit_set import QubitSet, QubitSetInput  # noqa: F401
+from braket.circuits.result_type import ResultType  # noqa: F401
+from braket.circuits.result_types import Amplitude, Probability, StateVector  # noqa: F401

src/braket/circuits/circuit.py

@@ -13,16 +13,19 @@
 
 from __future__ import annotations
 
-from typing import Callable, Dict, Iterable, TypeVar
+from typing import Callable, Dict, Iterable, List, TypeVar
 
 from braket.circuits.ascii_circuit_diagram import AsciiCircuitDiagram
 from braket.circuits.instruction import Instruction
 from braket.circuits.moments import Moments
 from braket.circuits.qubit import QubitInput
 from braket.circuits.qubit_set import QubitSet, QubitSetInput
+from braket.circuits.result_type import ResultType
 from braket.ir.jaqcd import Program
 
-SubroutineReturn = TypeVar("SubroutineReturn", Iterable[Instruction], Instruction)
+SubroutineReturn = TypeVar(
+    "SubroutineReturn", Iterable[Instruction], Instruction, ResultType, Iterable[ResultType]
+)
 SubroutineCallable = TypeVar("SubroutineCallable", bound=Callable[..., SubroutineReturn])
 AddableTypes = TypeVar("AddableTypes", SubroutineReturn, SubroutineCallable)
 
@@ -32,7 +35,14 @@
 class Circuit:
     """
     A representation of a quantum circuit that contains the instructions to be performed on a
-    quantum device. See :mod:`braket.circuits.gates` module for all of the supported instructions.
+    quantum device and the requested result types.
+
+    See :mod:`braket.circuits.gates` module for all of the supported instructions.
+
+    See :mod:`braket.circuits.result_types` module for all of the supported result types.
+
+    `AddableTypes` are `Instruction`, iterable of `Instruction`, `ResultType`,
+    iterable of `ResultType`, or `SubroutineCallable`
     """
 
     @classmethod
@@ -42,8 +52,8 @@ def register_subroutine(cls, func: SubroutineCallable) -> None:
         is the name of `func`.
 
         Args:
-            func (Callable[..., Union[Instruction, Iterable[Instruction]]]): The function of the
-                subroutine to add to the class.
+            func (Callable[..., Union[Instruction, Iterable[Instruction], ResultType,
+                Iterable[ResultType]]): The function of the subroutine to add to the class.
 
         Examples:
             >>> def h_on_all(target):
@@ -73,8 +83,8 @@ def method_from_subroutine(self, *args, **kwargs) -> SubroutineReturn:
     def __init__(self, addable: AddableTypes = None, *args, **kwargs):
         """
         Args:
-            addable (Instruction, iterable of `Instruction`, or `SubroutineCallable`, optional): The
-                instruction-like item(s) to add to self. Default = None.
+            addable (AddableTypes): The item(s) to add to self.
+                Default = None.
             *args: Variable length argument list. Supports any arguments that `add()` offers.
             **kwargs: Arbitrary keyword arguments. Supports any keyword arguments that `add()`
                 offers.
@@ -85,15 +95,18 @@ def __init__(self, addable: AddableTypes = None, *args, **kwargs):
         Examples:
             >>> circ = Circuit([Instruction(Gate.H(), 4), Instruction(Gate.CNot(), [4, 5])])
             >>> circ = Circuit().h(0).cnot(0, 1)
+            >>> circ = Circuit().h(0).cnot(0, 1).probability([0, 1])
 
             >>> @circuit.subroutine(register=True)
             >>> def bell_pair(target):
             ...     return Circ().h(target[0]).cnot(target[0:2])
             ...
             >>> circ = Circuit(bell_pair, [4,5])
             >>> circ = Circuit().bell_pair([4,5])
+
         """
         self._moments: Moments = Moments()
+        self._result_types: Iterable[ResultType] = []
 
         if addable is not None:
             self.add(addable, *args, **kwargs)
@@ -108,6 +121,11 @@ def instructions(self) -> Iterable[Instruction]:
         """Iterable[Instruction]: Get an `iterable` of instructions in the circuit."""
         return self._moments.values()
 
+    @property
+    def result_types(self) -> List[ResultType]:
+        """List[ResultType]: Get a list of requested result types in the circuit."""
+        return self._result_types
+
     @property
     def moments(self) -> Moments:
         """Moments: Get the `moments` for this circuit."""
@@ -123,6 +141,74 @@ def qubits(self) -> QubitSet:
         """QubitSet: Get a copy of the qubits for this circuit."""
         return QubitSet(self._moments.qubits)
 
+    def add_result_type(
+        self,
+        result_type: ResultType,
+        target: QubitSetInput = None,
+        target_mapping: Dict[QubitInput, QubitInput] = {},
+    ) -> Circuit:
+        """
+        Add a requested result type to `self`, returns `self` for chaining ability.
+
+        Args:
+            result_type (ResultType): `ResultType` to add into `self`.
+            target (int, Qubit, or iterable of int / Qubit, optional): Target qubits for the
+                `result_type`.
+                Default = None.
+            target_mapping (dictionary[int or Qubit, int or Qubit], optional): A dictionary of
+                qubit mappings to apply to the `result_type.target`. Key is the qubit in
+                `result_type.target` and the value is what the key will be changed to. Default = {}.
+
+
+        Note: target and target_mapping will only be applied to those requested result types with
+        the attribute `target`. The result_type will be appended to the end of the list of
+        `circuit.result_types` only if it does not already exist in `circuit.result_types`
+
+        Returns:
+            Circuit: self
+
+        Raises:
+            TypeError: If both `target_mapping` and `target` are supplied.
+
+        Examples:
+            >>> result_type = ResultType.Probability(target=[0, 1])
+            >>> circ = Circuit().add_result_type(result_type)
+            >>> print(circ.result_types[0])
+            Probability(target=QubitSet([Qubit(0), Qubit(1)]))
+
+            >>> result_type = ResultType.Probability(target=[0, 1])
+            >>> circ = Circuit().add_result_type(result_type, target_mapping={0: 10, 1: 11})
+            >>> print(circ.result_types[0])
+            Probability(target=QubitSet([Qubit(10), Qubit(11)]))
+
+            >>> result_type = ResultType.Probability(target=[0, 1])
+            >>> circ = Circuit().add_result_type(result_type, target=[10, 11])
+            >>> print(circ.result_types[0])
+            Probability(target=QubitSet([Qubit(10), Qubit(11)]))
+
+            >>> result_type = ResultType.StateVector()
+            >>> circ = Circuit().add_result_type(result_type)
+            >>> print(circ.result_types[0])
+            StateVector()
+        """
+        if target_mapping and target is not None:
+            raise TypeError("Only one of 'target_mapping' or 'target' can be supplied.")
+
+        if not target_mapping and not target:
+            # Nothing has been supplied, add result_type
+            result_type_to_add = result_type
+        elif target_mapping:
+            # Target mapping has been supplied, copy result_type
+            result_type_to_add = result_type.copy(target_mapping=target_mapping)
+        else:
+            # ResultType with target
+            result_type_to_add = result_type.copy(target=target)
+
+        if result_type_to_add not in self._result_types:
+            self._result_types.append(result_type_to_add)
+
+        return self
+
     def add_instruction(
         self,
         instruction: Instruction,
@@ -198,8 +284,7 @@ def add_circuit(
         target_mapping: Dict[QubitInput, QubitInput] = {},
     ) -> Circuit:
         """
-        Add a `circuit` to self, returns self for chaining ability. This is a composite form of
-        `add_instruction()` since it adds all of the instructions of `circuit` to this circuit.
+        Add a `circuit` to self, returns self for chaining ability.
 
         Args:
             circuit (Circuit): Circuit to add into self.
@@ -224,6 +309,10 @@ def add_circuit(
             can be resource-intensive. Use `target_mapping` to use a linear runtime to remap
             the qubits.
 
+            Requested result types of the circuit that will be added will be appended to the end
+            of the list for the existing requested result types. A result type to be added that is
+            equivalent to an existing requested result type will not be added.
+
         Examples:
             >>> widget = Circuit().h(0).cnot([0, 1])
             >>> circ = Circuit().add_circuit(widget)
@@ -256,18 +345,21 @@ def add_circuit(
         for instruction in circuit.instructions:
             self.add_instruction(instruction, target_mapping=target_mapping)
 
+        for result_type in circuit.result_types:
+            self.add_result_type(result_type, target_mapping=target_mapping)
+
         return self
 
     def add(self, addable: AddableTypes, *args, **kwargs) -> Circuit:
         """
-        Generic add method for adding instruction-like item(s) to self. Any arguments that
-        `add_circuit()` and / or `add_instruction()` supports are supported by this method.
-        If adding a subroutine, check with that subroutines documentation to determine what input it
+        Generic add method for adding item(s) to self. Any arguments that
+        `add_circuit()` and / or `add_instruction()` and / or `add_result_type`
+        supports are supported by this method. If adding a subroutine,
+        check with that subroutines documentation to determine what input it
         allows.
 
         Args:
-            addable (Instruction, iterable of Instruction, or SubroutineCallable, optional): The
-                instruction-like item(s) to add to self. Default = None.
+            addable (AddableTypes): The item(s) to add to self. Default = None.
             *args: Variable length argument list.
             **kwargs: Arbitrary keyword arguments.
 
@@ -282,8 +374,11 @@ def add(self, addable: AddableTypes, *args, **kwargs) -> Circuit:
 
             `add_instruction()`
 
+            `add_result_type()`
+
         Examples:
             >>> circ = Circuit().add([Instruction(Gate.H(), 4), Instruction(Gate.CNot(), [4, 5])])
+            >>> circ = Circuit().add([ResultType.StateVector()])
 
             >>> circ = Circuit().h(4).cnot([4, 5])
 
@@ -304,6 +399,8 @@ def _flatten(addable):
         for item in _flatten(addable):
             if isinstance(item, Instruction):
                 self.add_instruction(item, *args, **kwargs)
+            elif isinstance(item, ResultType):
+                self.add_result_type(item, *args, **kwargs)
             elif isinstance(item, Circuit):
                 self.add_circuit(item, *args, **kwargs)
             elif callable(item):
@@ -335,7 +432,8 @@ def to_ir(self) -> Program:
             (Program): An AWS quantum circuit description program in JSON format.
         """
         ir_instructions = [instr.to_ir() for instr in self.instructions]
-        return Program(instructions=ir_instructions)
+        ir_results = [result_type.to_ir() for result_type in self.result_types]
+        return Program(instructions=ir_instructions, results=ir_results)
 
     def _copy(self) -> Circuit:
         """
@@ -344,7 +442,9 @@ def _copy(self) -> Circuit:
         Returns:
             Circuit: A shallow copy of the circuit.
         """
-        return Circuit().add(self.instructions)
+        copy = Circuit().add(self.instructions)
+        copy.add(self.result_types)
+        return copy
 
     def __iadd__(self, addable: AddableTypes) -> Circuit:
         return self.add(addable)
@@ -354,21 +454,30 @@ def __add__(self, addable: AddableTypes) -> Circuit:
         new.add(addable)
         return new
 
-    def __repr__(self):
-        return f"Circuit('instructions': {list(self.instructions)})"
+    def __repr__(self) -> str:
+        if not self.result_types:
+            return f"Circuit('instructions': {list(self.instructions)})"
+        else:
+            return (
+                f"Circuit('instructions': {list(self.instructions)}"
+                + f"result_types': {self.result_types})"
+            )
 
     def __str__(self):
         return self.diagram(AsciiCircuitDiagram)
 
     def __eq__(self, other):
         if isinstance(other, Circuit):
-            return list(self.instructions) == list(other.instructions)
+            return (
+                list(self.instructions) == list(other.instructions)
+                and self.result_types == self.result_types
+            )
         return NotImplemented
 
 
 def subroutine(register=False):
     """
-    Subroutine is a function that returns instructions or circuits.
+    Subroutine is a function that returns instructions, result types, or circuits.
 
     Args:
         register (bool, optional): If `True`, adds this subroutine into the `Circuit` class.

src/braket/circuits/gate.py

@@ -17,9 +17,6 @@
 from braket.circuits.operator import Operator
 from braket.circuits.qubit_set import QubitSet
 
-# TODO: Add parameters support
-# TODO: Add printing / visualization capabilities
-
 
 class Gate(Operator):
     """
@@ -131,6 +128,6 @@ def register_gate(cls, gate: "Gate"):
         """Register a gate implementation by adding it into the Gate class.
 
         Args:
-            gate (Gate): Gate instance to register.
+            gate (Gate): Gate class to register.
         """
         setattr(cls, gate.__name__, gate)

src/braket/circuits/gates.py

@@ -23,8 +23,6 @@
 from braket.circuits.qubit import QubitInput
 from braket.circuits.qubit_set import QubitSet, QubitSetInput
 
-# TODO: look into adding angle to diagrams
-
 """
 To add a new gate:
     1. Implement the class and extend `Gate`

src/braket/circuits/operator.py

@@ -29,5 +29,6 @@ def to_ir(self, *args, **kwargs):
         If the operator is passed in a request, this method is called before it is passed.
 
         Args:
-            target (QubitSet): Target qubits that the operator is applied to.
+            *args: Positional arguments
+            **kwargs: Keyword arguments
         """