fix: Convert amplitude result to a complex number (#144)

src/braket/aws/aws_quantum_task.py

@@ -426,6 +426,7 @@ def _format_result(result):
 
 @_format_result.register
 def _(result: GateModelTaskResult) -> GateModelQuantumTaskResult:
+    GateModelQuantumTaskResult.cast_result_types(result)
     return GateModelQuantumTaskResult.from_object(result)
 
 

src/braket/tasks/gate_model_quantum_task_result.py

@@ -217,16 +217,7 @@ def from_string(result: str) -> GateModelQuantumTaskResult:
                 in the result dict
         """
         obj = GateModelTaskResult.parse_raw(result)
-        if obj.resultTypes:
-            for result_type in obj.resultTypes:
-                type = result_type.type.type
-                if type == "probability":
-                    result_type.value = np.array(result_type.value)
-                elif type == "statevector":
-                    result_type.value = np.array([complex(*value) for value in result_type.value])
-                elif type == "amplitude":
-                    for state in result_type.value:
-                        result_type.value[state] = complex(*result_type.value[state])
+        GateModelQuantumTaskResult.cast_result_types(obj)
         return GateModelQuantumTaskResult._from_object_internal(obj)
 
     @classmethod
@@ -303,6 +294,26 @@ def _from_dict_internal_simulator_only(cls, result: GateModelTaskResult):
             values=values,
         )
 
+    @staticmethod
+    def cast_result_types(gate_model_task_result: GateModelTaskResult) -> None:
+        """
+        Casts the result types to the types expected by the SDK.
+
+        Args:
+            gate_model_task_result (GateModelTaskResult): GateModelTaskResult representing the
+                results.
+        """
+        if gate_model_task_result.resultTypes:
+            for result_type in gate_model_task_result.resultTypes:
+                type = result_type.type.type
+                if type == "probability":
+                    result_type.value = np.array(result_type.value)
+                elif type == "statevector":
+                    result_type.value = np.array([complex(*value) for value in result_type.value])
+                elif type == "amplitude":
+                    for state in result_type.value:
+                        result_type.value[state] = complex(*result_type.value[state])
+
     @staticmethod
     def _calculate_result_types(
         ir_string: str, measurements: np.ndarray, measured_qubits: List[int]

test/integ_tests/gate_model_device_testing_utils.py

@@ -51,27 +51,31 @@ def no_result_types_bell_pair_testing(device: Device, run_kwargs: Dict[str, Any]
     assert len(result.measurements) == shots
 
 
-def result_types_zero_shots_bell_pair_testing(device: Device, run_kwargs: Dict[str, Any]):
+def result_types_zero_shots_bell_pair_testing(
+    device: Device, include_state_vector: bool, run_kwargs: Dict[str, Any]
+):
     circuit = (
         Circuit()
         .h(0)
         .cnot(0, 1)
         .expectation(observable=Observable.H() @ Observable.X(), target=[0, 1])
-        .state_vector()
         .amplitude(["01", "10", "00", "11"])
     )
+    if include_state_vector:
+        circuit.state_vector()
     result = device.run(circuit, **run_kwargs).result()
-    assert len(result.result_types) == 3
+    assert len(result.result_types) == 3 if include_state_vector else 2
     assert np.allclose(
         result.get_value_by_result_type(
             ResultType.Expectation(observable=Observable.H() @ Observable.X(), target=[0, 1])
         ),
         1 / np.sqrt(2),
     )
-    assert np.allclose(
-        result.get_value_by_result_type(ResultType.StateVector()),
-        np.array([1, 0, 0, 1]) / np.sqrt(2),
-    )
+    if include_state_vector:
+        assert np.allclose(
+            result.get_value_by_result_type(ResultType.StateVector()),
+            np.array([1, 0, 0, 1]) / np.sqrt(2),
+        )
     assert result.get_value_by_result_type(ResultType.Amplitude(["01", "10", "00", "11"])) == {
         "01": 0j,
         "10": 0j,

test/integ_tests/test_local_braket_simulator.py

@@ -44,7 +44,7 @@ def test_qubit_ordering():
 
 
 def test_result_types_no_shots():
-    result_types_zero_shots_bell_pair_testing(DEVICE, {"shots": 0})
+    result_types_zero_shots_bell_pair_testing(DEVICE, True, {"shots": 0})
 
 
 def test_result_types_nonzero_shots_bell_pair():

test/integ_tests/test_simulator_quantum_task.py

@@ -26,6 +26,7 @@
     result_types_tensor_z_h_y_testing,
     result_types_tensor_z_hermitian_testing,
     result_types_tensor_z_z_testing,
+    result_types_zero_shots_bell_pair_testing,
 )
 
 from braket.aws import AwsDevice
@@ -48,6 +49,14 @@ def test_qubit_ordering(simulator_arn, aws_session, s3_destination_folder):
     qubit_ordering_testing(device, {"shots": SHOTS, "s3_destination_folder": s3_destination_folder})
 
 
+@pytest.mark.parametrize("simulator_arn", [SIMULATOR_ARN])
+def test_result_types_no_shots(simulator_arn, aws_session, s3_destination_folder):
+    device = AwsDevice(simulator_arn, aws_session)
+    result_types_zero_shots_bell_pair_testing(
+        device, False, {"shots": 0, "s3_destination_folder": s3_destination_folder}
+    )
+
+
 @pytest.mark.parametrize("simulator_arn", [SIMULATOR_ARN])
 def test_result_types_nonzero_shots_bell_pair(simulator_arn, aws_session, s3_destination_folder):
     device = AwsDevice(simulator_arn, aws_session)
@@ -74,39 +83,39 @@ def test_result_types_bell_pair_marginal_probability(
     )
 
 
-@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS)])
+@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS), (SIMULATOR_ARN, 0)])
 def test_result_types_tensor_x_y(simulator_arn, shots, aws_session, s3_destination_folder):
     device = AwsDevice(simulator_arn, aws_session)
     result_types_tensor_x_y_testing(
         device, {"shots": shots, "s3_destination_folder": s3_destination_folder}
     )
 
 
-@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS)])
+@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS), (SIMULATOR_ARN, 0)])
 def test_result_types_tensor_z_h_y(simulator_arn, shots, aws_session, s3_destination_folder):
     device = AwsDevice(simulator_arn, aws_session)
     result_types_tensor_z_h_y_testing(
         device, {"shots": shots, "s3_destination_folder": s3_destination_folder}
     )
 
 
-@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS)])
+@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS), (SIMULATOR_ARN, 0)])
 def test_result_types_hermitian(simulator_arn, shots, aws_session, s3_destination_folder):
     device = AwsDevice(simulator_arn, aws_session)
     result_types_hermitian_testing(
         device, {"shots": shots, "s3_destination_folder": s3_destination_folder}
     )
 
 
-@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS)])
+@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS), (SIMULATOR_ARN, 0)])
 def test_result_types_tensor_z_z(simulator_arn, shots, aws_session, s3_destination_folder):
     device = AwsDevice(simulator_arn, aws_session)
     result_types_tensor_z_z_testing(
         device, {"shots": shots, "s3_destination_folder": s3_destination_folder}
     )
 
 
-@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS)])
+@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS), (SIMULATOR_ARN, 0)])
 def test_result_types_tensor_hermitian_hermitian(
     simulator_arn, shots, aws_session, s3_destination_folder
 ):
@@ -116,15 +125,15 @@ def test_result_types_tensor_hermitian_hermitian(
     )
 
 
-@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS)])
+@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS), (SIMULATOR_ARN, 0)])
 def test_result_types_tensor_y_hermitian(simulator_arn, shots, aws_session, s3_destination_folder):
     device = AwsDevice(simulator_arn, aws_session)
     result_types_tensor_y_hermitian_testing(
         device, {"shots": shots, "s3_destination_folder": s3_destination_folder}
     )
 
 
-@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS)])
+@pytest.mark.parametrize("simulator_arn,shots", [(SIMULATOR_ARN, SHOTS), (SIMULATOR_ARN, 0)])
 def test_result_types_tensor_z_hermitian(simulator_arn, shots, aws_session, s3_destination_folder):
     device = AwsDevice(simulator_arn, aws_session)
     result_types_tensor_z_hermitian_testing(

test/unit_tests/braket/aws/common_test_utils.py

@@ -45,6 +45,44 @@ class MockS3:
         }
     )
 
+    MOCK_S3_RESULT_GATE_MODEL_WITH_RESULT_TYPES = json.dumps(
+        {
+            "braketSchemaHeader": {
+                "name": "braket.task_result.gate_model_task_result",
+                "version": "1",
+            },
+            "measurements": [[0, 0], [0, 0], [0, 0], [1, 1]],
+            "measuredQubits": [0, 1],
+            "resultTypes": [
+                {
+                    "type": {"observable": ["h", "x"], "targets": [0, 1], "type": "expectation"},
+                    "value": 0.7071067811865474,
+                },
+                {
+                    "type": {"states": ["01", "10", "00", "11"], "type": "amplitude"},
+                    "value": {
+                        "01": [0.0, 0.0],
+                        "10": [0.0, 0.0],
+                        "00": [0.7071067811865475, 0.0],
+                        "11": [0.7071067811865475, 0.0],
+                    },
+                },
+            ],
+            "taskMetadata": {
+                "braketSchemaHeader": {"name": "braket.task_result.task_metadata", "version": "1"},
+                "id": "task_arn",
+                "shots": 100,
+                "deviceId": "default",
+            },
+            "additionalMetadata": {
+                "action": {
+                    "braketSchemaHeader": {"name": "braket.ir.jaqcd.program", "version": "1"},
+                    "instructions": [{"control": 0, "target": 1, "type": "cnot"}],
+                },
+            },
+        }
+    )
+
     MOCK_S3_RESULT_ANNEALING = json.dumps(
         {
             "braketSchemaHeader": {

test/unit_tests/braket/aws/test_aws_quantum_task.py

@@ -190,9 +190,13 @@ def test_result_annealing(annealing_task):
     )
 
 
-def test_result_is_cached(circuit_task):
+@pytest.mark.parametrize(
+    "result_string",
+    [MockS3.MOCK_S3_RESULT_GATE_MODEL, MockS3.MOCK_S3_RESULT_GATE_MODEL_WITH_RESULT_TYPES],
+)
+def test_result_is_cached(circuit_task, result_string):
     _mock_metadata(circuit_task._aws_session, "COMPLETED")
-    _mock_s3(circuit_task._aws_session, MockS3.MOCK_S3_RESULT_GATE_MODEL)
+    _mock_s3(circuit_task._aws_session, result_string)
     circuit_task.result()
 
     _mock_s3(circuit_task._aws_session, "")