Add results to jacqd program (#12)

CODEOWNERS

@@ -4,4 +4,4 @@
 # the repo. Unless a later match takes precedence, these accounts
 # will be requested forreview when someone opens a pull request.
 *   @floralph
-
+*   @avawang1

README.md

@@ -19,11 +19,11 @@ Amazon Braket Python IR is an open source library that contains all the intermed
  ```
 
 ## Usage
-There is currently only one kind of IR and that is the JsonAwsQuantumCircuitDescription, jaqcd. See below for it's usage.
+There are currently two types of IR, including jacqd (JsonAwsQuantumCircuitDescription) and annealing. See below for their usage.
 
 **Serializing python structures**
 ```python
-from braket.ir.jaqcd import CNot, H, Program
+from braket.ir.jaqcd import CNot, H, Program, Expectation
 from braket.ir.annealing import Problem, ProblemType
 
 program = Program(instructions=[H(target=0), CNot(control=0, target=1)])
@@ -41,6 +41,37 @@ print(program.json(indent=2))
       "target": 1,
       "type": "cnot"
     }
+  ],
+  "results": null
+}
+"""
+
+program = Program(instructions=[H(target=0), CNot(control=0, target=1)], results=[Expectation(targets=[1], observable=['x'])])
+print(program.json(indent=2))
+
+"""
+{
+  "instructions": [
+    {
+      "target": 0,
+      "type": "h"
+    },
+    {
+      "control": 0,
+      "target": 1,
+      "type": "cnot"
+    }
+  ],
+  "results": [
+    {
+      "observable": [
+        "x"
+      ],
+      "targets": [
+        1
+      ],
+      "type": "expectation"
+    }
   ]
 }
 """
@@ -74,6 +105,17 @@ jaqcd_string = """
       "target": 1,
       "type": "cnot"
     }
+  ],
+  "results": [
+    {
+      "observable": [
+        "x"
+      ],
+      "targets": [
+        1
+      ],
+      "type": "expectation"
+    }
   ]
 }
 """
@@ -82,7 +124,7 @@ program = Program.parse_raw(jaqcd_string)
 print(program)
 
 """
-instructions=[H(target=0, type=<Type.h: 'h'>), CNot(control=0, target=1, type=<Type.cnot: 'cnot'>)]
+instructions=[H(target=0, type=<Type.h: 'h'>), CNot(control=0, target=1, type=<Type.cnot: 'cnot'>)] results=[Expectation(observable=['x'], targets=[1], type=<Type.expectation: 'expectation'>)]
 """
 
 annealing_string = """

setup.py

@@ -15,7 +15,7 @@
 
 setup(
     name="braket-ir",
-    version="0.2.2",
+    version="0.2.3",
     license="Apache License 2.0",
     python_requires=">= 3.7",
     packages=find_namespace_packages(where="src", exclude=("test",)),

src/braket/ir/jaqcd/__init__.py

@@ -46,3 +46,11 @@
     Z,
 )
 from braket.ir.jaqcd.program import Program  # noqa: F401
+from braket.ir.jaqcd.results import (  # noqa: F401
+    Amplitude,
+    Expectation,
+    Probability,
+    Sample,
+    StateVector,
+    Variance,
+)

src/braket/ir/jaqcd/program.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from typing import List, Union
+from typing import List, Optional, Union
 
 from braket.ir.jaqcd.instructions import (
     CY,
@@ -47,6 +47,14 @@
     Y,
     Z,
 )
+from braket.ir.jaqcd.results import (
+    Amplitude,
+    Expectation,
+    Probability,
+    Sample,
+    StateVector,
+    Variance,
+)
 from pydantic import BaseModel
 
 
@@ -97,3 +105,6 @@ class Program(BaseModel):
             ZZ,
         ]
     ]
+    results: Optional[
+        List[Union[Amplitude, Expectation, Probability, Sample, StateVector, Variance]]
+    ]

src/braket/ir/jaqcd/results.py

@@ -0,0 +1,159 @@
+# Copyright 2019-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+
+from enum import Enum
+
+from braket.ir.jaqcd.shared_models import MultiState, Observable, OptionalMultiTarget
+from pydantic import BaseModel
+
+
+class Expectation(OptionalMultiTarget, Observable):
+    """
+    Expectation of specified targets and observable as requested result.
+    If no targets are specified, the observable must only operate on 1 qubit and it
+    will be applied to all qubits in parallel. Otherwise, the number of specified targets
+    must be equivalent to the number of qubits the observable can be applied to.
+
+    Attributes:
+        type (str): The result type. default = "expectation". (type) is optional.
+            This should be unique among all result types.
+        targets (Optional[List[int]]): The target qubits. This is a list of int >= 0.
+        observable (List[Union[str, List[List[List[float]]]]): A list with at least
+            one item and items are strings matching the observable regex
+            or a two dimensional hermitian matrix with complex entries.
+            Each complex number is represented using a List[float] of size 2, with
+            element[0] being the real part and element[1] imaginary.
+            inf, -inf, and NaN are not allowable inputs for the element.
+
+    Examples:
+        >>> Expectation(targets=[1], observable=["x"])
+    """
+
+    class Type(str, Enum):
+        expectation = "expectation"
+
+    type = Type.expectation
+
+
+class Sample(OptionalMultiTarget, Observable):
+    """
+    Sample for specified targets and observable as requested result.
+    If no targets are specified, the observable must only operate on 1 qubit and it
+    will be applied to all qubits in parallel. Otherwise, the number of specified targets
+    must be equivalent to the number of qubits the observable can be applied to.
+
+    Attributes:
+        type (str): The result type. default = "sample". (type) is optional.
+            This should be unique among all result types.
+        targets (Optional[List[int]]): The target qubits. This is a list of int >= 0.
+        observable (List[Union[str, List[List[List[float]]]]): A list with at least
+            one item and items are strings matching the observable regex
+            or a two dimensional hermitian matrix with complex entries.
+            Each complex number is represented using a List[float] of size 2, with
+            element[0] being the real part and element[1] imaginary.
+            inf, -inf, and NaN are not allowable inputs for the element.
+
+    Examples:
+        >>> Sample(targets=[1], observable=["x"])
+    """
+
+    class Type(str, Enum):
+        sample = "sample"
+
+    type = Type.sample
+
+
+class Variance(OptionalMultiTarget, Observable):
+    """
+    Variance of specified targets and observables as requested result.
+    If no targets are specified, the observable must only operate on 1 qubit and it
+    will be applied to all qubits in parallel. Otherwise, the number of specified targets
+    must be equivalent to the number of qubits the observable can be applied to.
+
+    Attributes:
+        type (str): The result type. default = "variance". (type) is optional.
+            This should be unique among all result types.
+        targets (List[int]): The target qubits. This is a list of int >= 0.
+        observable (List[Union[str, List[List[List[float]]]]): A list with at least
+            one item and items are strings matching the observable regex
+            or a two dimensional hermitian matrix with complex entries.
+            Each complex number is represented using a List[float] of size 2, with
+            element[0] being the real part and element[1] imaginary.
+            inf, -inf, and NaN are not allowable inputs for the element.
+
+    Examples:
+        >>> Variance(targets=[1], observable=["x"])
+    """
+
+    class Type(str, Enum):
+        variance = "variance"
+
+    type = Type.variance
+
+
+class StateVector(BaseModel):
+    """
+    The full state vector as requested result.
+
+    Attributes:
+        type (str): The result type. default = "statevector". (type) is optional.
+            This should be unique among all result types.
+
+    Examples:
+        >>> StateVector()
+    """
+
+    class Type(str, Enum):
+        statevector = "statevector"
+
+    type = Type.statevector
+
+
+class Amplitude(MultiState):
+    """
+    Amplitudes of specified states as requested result.
+
+    Attributes:
+        type (str): The result type. default = "amplitude". (type) is optional.
+            This should be unique among all result types.
+        states (List[string]): Variable length list with with all strings
+            matching the state regex
+
+    Examples:
+        >>> Amplitude(states=["01", "10"])
+    """
+
+    class Type(str, Enum):
+        amplitude = "amplitude"
+
+    type = Type.amplitude
+
+
+class Probability(OptionalMultiTarget):
+    """
+    Probability of all states if no targets are specified or the marginal probability
+        of a restricted set of states if only a subset of all qubits are specified as targets
+
+    Attributes:
+        type (str): The result type. default = "probability". (type) is optional.
+            This should be unique among all result types.
+        targets (Optional[List[int]]): The target qubits. This is a list of int >= 0.
+
+    Examples:
+        >>> Probability(targets=[1, 2])
+    """
+
+    class Type(str, Enum):
+        probability = "probability"
+
+    type = Type.probability

src/braket/ir/jaqcd/shared_models.py

@@ -11,7 +11,9 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from pydantic import BaseModel, confloat, conint, conlist
+from typing import Optional, Union
+
+from pydantic import BaseModel, confloat, conint, conlist, constr
 
 
 class SingleTarget(BaseModel):
@@ -47,7 +49,7 @@ class MultiTarget(BaseModel):
     Variable length target indices.
 
     Attributes:
-        targets (List[int]): A list with at least two items and all items are int >= 0.
+        targets (List[int]): A list with items that are all int >= 0.
 
     Examples:
         >>> MultiTarget(targets=[0, 1])
@@ -56,6 +58,20 @@ class MultiTarget(BaseModel):
     targets: conlist(conint(ge=0), min_items=1)
 
 
+class OptionalMultiTarget(BaseModel):
+    """
+    Optional variable length target indices
+
+    Attributes:
+        targets (Optional[List[int]]): A list with items that are all int >= 0.
+
+    Examples:
+        >>> OptionalMultiTarget(targets=[0, 1])
+    """
+
+    targets: Optional[conlist(conint(ge=0), min_items=1)]
+
+
 class MultiControl(BaseModel):
     """
     Variable length control indices.
@@ -134,3 +150,51 @@ class TwoDimensionalMatrix(BaseModel):
         ),
         min_items=1,
     )
+
+
+class Observable(BaseModel):
+    """
+    An observable. If given list is more than one element, this is the tensor product
+    of each operator in the list.
+
+    Attributes:
+        observable (List[Union[str, List[List[List[float]]]]): A list with at least
+            one item and items are strings matching the observable regex
+            or a two dimensional hermitian matrix with complex entries.
+            Each complex number is represented using a List[float] of size 2, with
+            element[0] being the real part and element[1] imaginary.
+            inf, -inf, and NaN are not allowable inputs for the element.
+
+    Examples:
+        >>> Observable(observable=["x"])
+        >>> Observable(observable=[[[0, 0], [1, 0]], [[1, 0], [0, 0]]])
+    """
+
+    observable: conlist(
+        Union[
+            constr(regex="(x|y|z|h|i)"),
+            conlist(
+                conlist(
+                    conlist(confloat(gt=float("-inf"), lt=float("inf")), min_items=2, max_items=2),
+                    min_items=2,
+                ),
+                min_items=2,
+            ),
+        ],
+        min_items=1,
+    )
+
+
+class MultiState(BaseModel):
+    """
+    A list of states in bitstring form.
+
+    Attributes:
+        states (List[string]): Variable length list with all strings matching the
+            state regex
+
+    Examples:
+        >>> lMultiState(states=["10", "10"])
+    """
+
+    states: conlist(constr(regex="^[01]+$", min_length=1), min_items=1)