Calculate result types for shots > 0

src/braket/circuits/observables.py

@@ -15,7 +15,7 @@
 import functools
 import itertools
 import math
-from typing import Dict, List, Tuple
+from typing import Dict, List, Tuple, Union
 
 import numpy as np
 from braket.circuits.gate import Gate
@@ -329,3 +329,42 @@ def _get_eigendecomposition(self) -> Dict[str, np.ndarray]:
 
 
 Observable.register_observable(Hermitian)
+
+
+def observable_from_ir(ir_observable: List[Union[str, List[List[List[float]]]]]) -> Observable:
+    """
+    Create an observable from the IR observable list. This can be a tensor product of
+    observables or a single observable.
+
+    Args:
+        ir_observable (List[Union[str, List[List[List[float]]]]]): observable as defined in IR
+
+    Return:
+        Observable: observable object
+    """
+    if len(ir_observable) == 1:
+        return _observable_from_ir_list_item(ir_observable[0])
+    else:
+        observable = TensorProduct([_observable_from_ir_list_item(obs) for obs in ir_observable])
+        return observable
+
+
+def _observable_from_ir_list_item(observable: Union[str, List[List[List[float]]]]) -> Observable:
+    if observable == "i":
+        return I()
+    elif observable == "h":
+        return H()
+    elif observable == "x":
+        return X()
+    elif observable == "y":
+        return Y()
+    elif observable == "z":
+        return Z()
+    else:
+        try:
+            matrix = np.array(
+                [[complex(element[0], element[1]) for element in row] for row in observable]
+            )
+            return Hermitian(matrix)
+        except Exception as e:
+            raise ValueError(f"Invalid observable specified: {observable} error: {e}")

src/braket/tasks/gate_model_quantum_task_result.py

@@ -15,10 +15,11 @@
 
 import json
 from dataclasses import dataclass
-from typing import Any, Counter, Dict, List
+from typing import Any, Counter, Dict, List, Optional
 
 import numpy as np
-from braket.circuits import ResultType
+from braket.circuits import Observable, ResultType, StandardObservable
+from braket.circuits.observables import observable_from_ir
 
 
 @dataclass
@@ -30,12 +31,16 @@ class GateModelQuantumTaskResult:
     Args:
         task_metadata (Dict[str, Any]): Dictionary of task metadata. task_metadata must have
             keys 'Id', 'Shots', 'Ir', and 'IrType'.
-        result_types (List[Dict[str, Any]], optional): List of dictionaries where each dictionary
+        result_types (List[Dict[str, Any]]): List of dictionaries where each dictionary
             has two keys: 'Type' (the result type in IR JSON form) and
             'Value' (the result value for this result type).
-            Default is None. Currently only available when shots = 0. TODO: update
-        values (List[Any], optional): The values for result types requested in the circuit.
-            Default is None. Currently only available when shots = 0. TODO: update
+            This can be an empty list if no result types are specified in the IR.
+            This is calculated from `measurements` and
+            the IR of the circuit program when `shots>0`.
+        values (List[Any]): The values for result types requested in the circuit.
+            This can be an empty list if no result types are specified in the IR.
+            This is calculated from `measurements` and
+            the IR of the circuit program when `shots>0`.
         measurements (numpy.ndarray, optional): 2d array - row is shot, column is qubit.
             Default is None. Only available when shots > 0. The qubits in `measurements`
             are the ones in `GateModelQuantumTaskResult.measured_qubits`.
@@ -63,8 +68,8 @@ class GateModelQuantumTaskResult:
     """
 
     task_metadata: Dict[str, Any]
-    result_types: List[Dict[str, str]] = None
-    values: List[Any] = None
+    result_types: List[Dict[str, str]]
+    values: List[Any]
     measurements: np.ndarray = None
     measured_qubits: List[int] = None
     measurement_counts: Counter = None
@@ -225,7 +230,6 @@ def _from_dict_internal(cls, result: Dict[str, Any]):
     @classmethod
     def _from_dict_internal_computational_basis_sampling(cls, result: Dict[str, Any]):
         task_metadata = result["TaskMetadata"]
-        measured_qubits = result.get("MeasuredQubits")
         if "Measurements" in result:
             measurements = np.asarray(result["Measurements"], dtype=int)
             m_counts = GateModelQuantumTaskResult.measurement_counts_from_measurements(measurements)
@@ -249,9 +253,20 @@ def _from_dict_internal_computational_basis_sampling(cls, result: Dict[str, Any]
             raise ValueError(
                 'One of "Measurements" or "MeasurementProbabilities" must be in the result dict'
             )
-        # TODO: add result types for shots > 0 after basis rotation instructions are added to IR
+        measured_qubits = result["MeasuredQubits"]
+        if len(measured_qubits) != measurements.shape[1]:
+            raise ValueError(
+                f"Measured qubits {measured_qubits} is not equivalent to number of qubits "
+                + f"{measurements.shape[1]} in measurements"
+            )
+        result_types = GateModelQuantumTaskResult._calculate_result_types(
+            result["TaskMetadata"]["Ir"], measurements, measured_qubits
+        )
+        values = [rt["Value"] for rt in result_types]
         return cls(
             task_metadata=task_metadata,
+            result_types=result_types,
+            values=values,
             measurements=measurements,
             measured_qubits=measured_qubits,
             measurement_counts=m_counts,
@@ -267,3 +282,133 @@ def _from_dict_internal_simulator_only(cls, result: Dict[str, Any]):
         result_types = result["ResultTypes"]
         values = [rt["Value"] for rt in result_types]
         return cls(task_metadata=task_metadata, result_types=result_types, values=values)
+
+    @staticmethod
+    def _calculate_result_types(
+        ir_string: str, measurements: np.ndarray, measured_qubits: List[int]
+    ) -> List[Dict[str, Any]]:
+        ir = json.loads(ir_string)
+        result_types = []
+        if not ir.get("results"):
+            return result_types
+        for result_type in ir["results"]:
+            ir_observable = result_type.get("observable")
+            observable = observable_from_ir(ir_observable) if ir_observable else None
+            targets = result_type.get("targets")
+            rt_type = result_type["type"]
+            if rt_type == "probability":
+                value = GateModelQuantumTaskResult._probability_from_measurements(
+                    measurements, measured_qubits, targets
+                )
+            elif rt_type == "sample":
+                if targets:
+                    value = GateModelQuantumTaskResult._samples_from_measurements(
+                        measurements, measured_qubits, observable, targets
+                    )
+                else:
+                    value = [
+                        GateModelQuantumTaskResult._samples_from_measurements(
+                            measurements, measured_qubits, observable, [i]
+                        )
+                        for i in measured_qubits
+                    ]
+            elif rt_type == "variance":
+                if targets:
+                    value = GateModelQuantumTaskResult._variance_from_measurements(
+                        measurements, measured_qubits, observable, targets
+                    )
+                else:
+                    value = [
+                        GateModelQuantumTaskResult._variance_from_measurements(
+                            measurements, measured_qubits, observable, [i]
+                        )
+                        for i in measured_qubits
+                    ]
+            elif rt_type == "expectation":
+                if targets:
+                    value = GateModelQuantumTaskResult._expectation_from_measurements(
+                        measurements, measured_qubits, observable, targets
+                    )
+                else:
+                    value = [
+                        GateModelQuantumTaskResult._expectation_from_measurements(
+                            measurements, measured_qubits, observable, [i]
+                        )
+                        for i in measured_qubits
+                    ]
+            else:
+                raise ValueError(f"Unknown result type {rt_type}")
+            result_types.append({"Type": result_type, "Value": value})
+        return result_types
+
+    @staticmethod
+    def _selected_measurements(
+        measurements: np.ndarray, measured_qubits: List[int], targets: Optional[List[int]]
+    ) -> np.ndarray:
+        if not (targets is None or targets == measured_qubits):
+            # Only some qubits targeted
+            columns = [measured_qubits.index(t) for t in targets]
+            measurements = measurements[:, columns]
+        return measurements
+
+    @staticmethod
+    def _probability_from_measurements(
+        measurements: np.ndarray, measured_qubits: List[int], targets: Optional[List[int]]
+    ) -> np.ndarray:
+        measurements = GateModelQuantumTaskResult._selected_measurements(
+            measurements, measured_qubits, targets
+        )
+        shots, num_measured_qubits = measurements.shape
+        # convert samples from a list of 0, 1 integers, to base 10 representation
+        unraveled_indices = [2] * num_measured_qubits
+        indices = np.ravel_multi_index(measurements.T, unraveled_indices)
+
+        # count the basis state occurrences, and construct the probability vector
+        basis_states, counts = np.unique(indices, return_counts=True)
+        probabilities = np.zeros([2 ** num_measured_qubits], dtype=np.float64)
+        probabilities[basis_states] = counts / shots
+        return probabilities
+
+    @staticmethod
+    def _variance_from_measurements(
+        measurements: np.ndarray,
+        measured_qubits: List[int],
+        observable: Observable,
+        targets: List[int],
+    ) -> float:
+        samples = GateModelQuantumTaskResult._samples_from_measurements(
+            measurements, measured_qubits, observable, targets
+        )
+        return np.var(samples)
+
+    @staticmethod
+    def _expectation_from_measurements(
+        measurements: np.ndarray,
+        measured_qubits: List[int],
+        observable: Observable,
+        targets: List[int],
+    ) -> float:
+        samples = GateModelQuantumTaskResult._samples_from_measurements(
+            measurements, measured_qubits, observable, targets
+        )
+        return np.mean(samples)
+
+    @staticmethod
+    def _samples_from_measurements(
+        measurements: np.ndarray,
+        measured_qubits: List[int],
+        observable: Observable,
+        targets: List[int],
+    ) -> np.ndarray:
+        measurements = GateModelQuantumTaskResult._selected_measurements(
+            measurements, measured_qubits, targets
+        )
+        shots, num_measured_qubits = measurements.shape
+        if isinstance(observable, StandardObservable):
+            # Process samples for observables with eigenvalues {1, -1}
+            return 1 - 2 * measurements.flatten()
+        # Replace the basis state in the computational basis with the correct eigenvalue.
+        # Extract only the columns of the basis samples required based on ``wires``.
+        unraveled_indices = [2] * num_measured_qubits
+        indices = np.ravel_multi_index(measurements.T, unraveled_indices)
+        return observable.eigenvalues[indices]

test/unit_tests/braket/aws/common_test_utils.py

@@ -175,8 +175,8 @@ class MockS3:
 
     MOCK_S3_RESULT_1 = json.dumps(
         {
-            "StateVector": {"00": [0.2, 0.2], "01": [0.3, 0.1], "10": [0.1, 0.3], "11": [0.2, 0.2]},
             "Measurements": [[0, 0], [0, 1], [0, 1], [0, 1]],
+            "MeasuredQubits": [0, 1],
             "TaskMetadata": {
                 "Id": "UUID_blah_1",
                 "Status": "COMPLETED",
@@ -189,8 +189,8 @@ class MockS3:
 
     MOCK_S3_RESULT_2 = json.dumps(
         {
-            "StateVector": {"00": [0.2, 0.2], "01": [0.3, 0.1], "10": [0.1, 0.3], "11": [0.2, 0.2]},
             "Measurements": [[0, 0], [0, 0], [0, 0], [1, 1]],
+            "MeasuredQubits": [0, 1],
             "TaskMetadata": {
                 "Id": "UUID_blah_2",
                 "Status": "COMPLETED",
@@ -215,6 +215,7 @@ class MockS3:
                 "Ir": "{}",
             },
             "MeasurementProbabilities": {"011000": 0.9999999999999982},
+            "MeasuredQubits": [0, 1],
         }
     )
 

test/unit_tests/braket/circuits/test_observables.py

@@ -16,6 +16,7 @@
 import numpy as np
 import pytest
 from braket.circuits import Gate, Observable
+from braket.circuits.observables import observable_from_ir
 from braket.circuits.quantum_operator_helpers import get_pauli_eigenvalues
 
 testdata = [
@@ -63,6 +64,29 @@ def test_gate_equality(testobject, gateobject, expected_ir, basis_rotation_gates
     assert np.allclose(testobject.eigenvalues, eigenvalues)
 
 
+@pytest.mark.parametrize(
+    "testobject,gateobject,expected_ir,basis_rotation_gates,eigenvalues", testdata
+)
+def test_basis_rotation_gates(
+    testobject, gateobject, expected_ir, basis_rotation_gates, eigenvalues
+):
+    assert testobject.basis_rotation_gates == basis_rotation_gates
+
+
+@pytest.mark.parametrize(
+    "testobject,gateobject,expected_ir,basis_rotation_gates,eigenvalues", testdata
+)
+def test_eigenvalues(testobject, gateobject, expected_ir, basis_rotation_gates, eigenvalues):
+    assert np.allclose(testobject.eigenvalues, eigenvalues)
+
+
+@pytest.mark.parametrize(
+    "testobject,gateobject,expected_ir,basis_rotation_gates,eigenvalues", testdata
+)
+def test_observable_from_ir(testobject, gateobject, expected_ir, basis_rotation_gates, eigenvalues):
+    assert testobject == observable_from_ir(expected_ir)
+
+
 # Hermitian
 
 
@@ -134,6 +158,18 @@ def test_hermitian_basis_rotation_gates(matrix, basis_rotation_matrix):
     assert expected_unitary.matrix_equivalence(actual_rotation_gates)
 
 
+@pytest.mark.xfail(raises=ValueError)
+def test_observable_from_ir_hermitian_value_error():
+    ir_observable = [[[[1.0, 0], [0, 1]], [[0.0, 1], [1, 0]]]]
+    observable_from_ir(ir_observable)
+
+
+def test_observable_from_ir_hermitian():
+    ir_observable = [[[[1, 0], [0, 0]], [[0, 0], [1, 0]]]]
+    actual_observable = observable_from_ir(ir_observable)
+    assert actual_observable == Observable.Hermitian(matrix=np.array([[1.0, 0.0], [0.0, 1.0]]))
+
+
 # TensorProduct
 
 
@@ -212,3 +248,14 @@ def test_tensor_product_eigenvalues(observable, eigenvalues):
 )
 def test_tensor_product_basis_rotation_gates(observable, basis_rotation_gates):
     assert observable.basis_rotation_gates == basis_rotation_gates
+
+
+def test_observable_from_ir_tensor_product():
+    expected_observable = Observable.TensorProduct([Observable.Z(), Observable.I(), Observable.X()])
+    actual_observable = observable_from_ir(["z", "i", "x"])
+    assert expected_observable == actual_observable
+
+
+@pytest.mark.xfail(raises=ValueError)
+def test_observable_from_ir_tensor_product_value_error():
+    observable_from_ir(["z", "i", "foo"])

test/unit_tests/braket/devices/test_local_simulator.py

@@ -25,6 +25,7 @@
 
 GATE_MODEL_RESULT = {
     "Measurements": [[0, 0], [0, 1], [0, 1], [0, 1]],
+    "MeasuredQubits": [0, 1],
     "TaskMetadata": {
         "Id": "UUID_blah_1",
         "Status": "COMPLETED",