Add from_dict to results (#58)

* This eliminates the need for local simulators to go through an extra
ser/de results step
* Updated LocalSimulator and BraketSimulator to use this new method
* Also improved formatting for a couple of files

src/braket/aws/aws_qpu.py

@@ -10,6 +10,7 @@
 # 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 typing import Any, Dict, Optional, Union
 
 import boto3

src/braket/devices/braket_simulator.py

@@ -10,8 +10,9 @@
 # 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 abc import ABC, abstractmethod
-from typing import Union
+from typing import Any, Dict, Union
 
 from braket.ir.annealing import Problem
 from braket.ir.jaqcd import Program
@@ -41,7 +42,7 @@ class BraketSimulator(ABC):
     # TODO: Update to use new simulate() method
 
     @abstractmethod
-    def run(self, ir: Union[Program, Problem], *args, **kwargs) -> str:
+    def run(self, ir: Union[Program, Problem], *args, **kwargs) -> Dict[str, Any]:
         """ Run the task specified by the given IR.
 
         Extra arguments will contain any additional information necessary to run the task,
@@ -51,9 +52,9 @@ def run(self, ir: Union[Program, Problem], *args, **kwargs) -> str:
             ir (Union[Program, Problem]): The IR representation of the program
 
         Returns:
-            str: A JSON string containing the results of the simulation.
+            Dict[str, Any]: A dict containing the results of the simulation.
             In order to work with braket-python-sdk, the format of the JSON dict should
             match that needed by GateModelQuantumTaskResult or AnnealingQuantumTaskResult
-            in the SDK, depending on the type of task.
+            from the SDK, depending on the type of task.
         """
         raise NotImplementedError()

src/braket/devices/local_simulator.py

@@ -10,6 +10,7 @@
 # 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 functools import singledispatch
 from typing import Set, Union
 
@@ -69,8 +70,8 @@ def run(
         result = _run_internal(task_specification, self._delegate, *args, **kwargs)
         return LocalQuantumTask(result)
 
-    @classmethod
-    def registered_backends(cls) -> Set[str]:
+    @staticmethod
+    def registered_backends() -> Set[str]:
         """ Gets the backends that have been registered as entry points
 
         Returns:
@@ -108,12 +109,12 @@ def _run_internal(task_specification, simulator: BraketSimulator, *args, **kwarg
 def _(circuit: Circuit, simulator: BraketSimulator, *args, **kwargs):
     program = circuit.to_ir()
     qubits = circuit.qubit_count
-    result_json = simulator.run(program, qubits, *args, **kwargs)
-    return GateModelQuantumTaskResult.from_string(result_json)
+    results_dict = simulator.run(program, qubits, *args, **kwargs)
+    return GateModelQuantumTaskResult.from_dict(results_dict)
 
 
 @_run_internal.register
 def _(problem: Problem, simulator: BraketSimulator, *args, **kwargs):
     ir = problem.to_ir()
-    result_json = simulator.run(ir, *args, *kwargs)
-    return AnnealingQuantumTaskResult.from_string(result_json)
+    results_dict = simulator.run(ir, *args, *kwargs)
+    return AnnealingQuantumTaskResult.from_dict(results_dict)

src/braket/tasks/annealing_quantum_task_result.py

@@ -91,6 +91,19 @@ def __eq__(self, other) -> bool:
             return (self.record_array == other.record_array).all() and self_fields == other_fields
         return NotImplemented
 
+    @staticmethod
+    def from_dict(result: Dict[str, Any]):
+        """
+        Create AnnealingQuantumTaskResult from dict
+
+        Args:
+            result (Dict[str, Any]): Results dict with AnnealingQuantumTaskResult attributes as keys
+
+        Returns:
+            AnnealingQuantumTaskResult: An AnnealingQuantumTaskResult based on the given dict
+        """
+        return AnnealingQuantumTaskResult._from_dict_internal(result)
+
     @staticmethod
     def from_string(result: str) -> AnnealingQuantumTaskResult:
         """
@@ -100,26 +113,29 @@ def from_string(result: str) -> AnnealingQuantumTaskResult:
             result (str): JSON object string
 
         Returns:
-            AnnealingQuantumTaskResult: An AnnealingQuantumTaskResult based on a string.
+            AnnealingQuantumTaskResult: An AnnealingQuantumTaskResult based on the given string
         """
-        json_obj = json.loads(result)
-        solutions = numpy.asarray(json_obj["Solutions"], dtype=int)
-        values = numpy.asarray(json_obj["Values"], dtype=float)
-        if json_obj["SolutionCounts"] is None:
+        return AnnealingQuantumTaskResult._from_dict_internal(json.loads(result))
+
+    @classmethod
+    def _from_dict_internal(cls, result: Dict[str, Any]):
+        solutions = numpy.asarray(result["Solutions"], dtype=int)
+        values = numpy.asarray(result["Values"], dtype=float)
+        if result["SolutionCounts"] is None:
             solution_counts = numpy.ones(len(solutions), dtype=int)
         else:
-            solution_counts = numpy.asarray(json_obj["SolutionCounts"], dtype=int)
-        record_array = AnnealingQuantumTaskResult.create_record_array(
+            solution_counts = numpy.asarray(result["SolutionCounts"], dtype=int)
+        record_array = AnnealingQuantumTaskResult._create_record_array(
             solutions, solution_counts, values
         )
-        variable_count = json_obj["VariableCount"]
-        problem_type = json_obj["ProblemType"]
-        task_metadata = json_obj["TaskMetadata"]
+        variable_count = result["VariableCount"]
+        problem_type = result["ProblemType"]
+        task_metadata = result["TaskMetadata"]
         additional_metadata = {}
-        for key in json_obj.keys():
+        for key in result.keys():
             if key.endswith("Metadata") and key != "TaskMetadata":
-                additional_metadata[key] = json_obj[key]
-        return AnnealingQuantumTaskResult(
+                additional_metadata[key] = result[key]
+        return cls(
             record_array=record_array,
             variable_count=variable_count,
             problem_type=problem_type,
@@ -128,7 +144,7 @@ def from_string(result: str) -> AnnealingQuantumTaskResult:
         )
 
     @staticmethod
-    def create_record_array(
+    def _create_record_array(
         solutions: numpy.ndarray, solution_counts: numpy.ndarray, values: numpy.ndarray
     ) -> numpy.recarray:
         """

src/braket/tasks/gate_model_quantum_task_result.py

@@ -130,39 +130,70 @@ def measurements_from_measurement_probabilities(
             measurements_list.extend(individual_measurement_list)
         return np.asarray(measurements_list, dtype=int)
 
+    @staticmethod
+    def from_dict(result: Dict[str, Any]):
+        """
+        Create GateModelQuantumTaskResult from dict.
+
+        Args:
+            result (Dict[str, Any]): Results dict with GateModelQuantumTaskResult attributes as keys
+
+        Returns:
+            GateModelQuantumTaskResult: A GateModelQuantumTaskResult based on the given dict
+
+        Raises:
+            ValueError: If neither "Measurements" nor "MeasurementProbabilities" is a key
+            in the result dict
+
+        Note:
+            For the "StateVector" key, the value should be of type Dict[str, complex];
+            each bitstring's amplitude is a Python complex number.
+        """
+        return GateModelQuantumTaskResult._from_dict_internal(result)
+
     @staticmethod
     def from_string(result: str) -> GateModelQuantumTaskResult:
         """
-        Create GateModelQuantumTaskResult from string
+        Create GateModelQuantumTaskResult from string.
 
         Args:
             result (str): JSON object string, with GateModelQuantumTaskResult attributes as keys.
 
         Returns:
-            GateModelQuantumTaskResult: A GateModelQuantumTaskResult based on a string
+            GateModelQuantumTaskResult: A GateModelQuantumTaskResult based on the given string
 
         Raises:
             ValueError: If neither "Measurements" nor "MeasurementProbabilities" is a key
             in the result dict
+
+        Note:
+            For the "StateVector" key, the value should be of type Dict[str, List[float, float]];
+            each bitstring's amplitude is represented by a two-element list, with first element
+            the real component and second element the imaginary component.
         """
         json_obj = json.loads(result)
-        task_metadata = json_obj["TaskMetadata"]
         state_vector = json_obj.get("StateVector", None)
         if state_vector:
             for state in state_vector:
                 state_vector[state] = complex(*state_vector[state])
-        if "Measurements" in json_obj:
-            measurements = np.asarray(json_obj["Measurements"], dtype=int)
+        return GateModelQuantumTaskResult._from_dict_internal(json_obj)
+
+    @classmethod
+    def _from_dict_internal(cls, result: Dict[str, Any]):
+        task_metadata = result["TaskMetadata"]
+        state_vector = result.get("StateVector", None)
+        if "Measurements" in result:
+            measurements = np.asarray(result["Measurements"], dtype=int)
             m_counts = GateModelQuantumTaskResult.measurement_counts_from_measurements(measurements)
             m_probs = GateModelQuantumTaskResult.measurement_probabilities_from_measurement_counts(
                 m_counts
             )
             measurements_copied_from_device = True
             m_counts_copied_from_device = False
             m_probabilities_copied_from_device = False
-        elif "MeasurementProbabilities" in json_obj:
+        elif "MeasurementProbabilities" in result:
             shots = task_metadata["Shots"]
-            m_probs = json_obj["MeasurementProbabilities"]
+            m_probs = result["MeasurementProbabilities"]
             measurements = GateModelQuantumTaskResult.measurements_from_measurement_probabilities(
                 m_probs, shots
             )
@@ -174,7 +205,7 @@ def from_string(result: str) -> GateModelQuantumTaskResult:
             raise ValueError(
                 'One of "Measurements" or "MeasurementProbabilities" must be in the results dict'
             )
-        return GateModelQuantumTaskResult(
+        return cls(
             state_vector=state_vector,
             task_metadata=task_metadata,
             measurements=measurements,

src/braket/tasks/local_quantum_task.py

@@ -10,6 +10,7 @@
 # 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.
+
 import asyncio
 import uuid
 from typing import Union

test/integ_tests/test_simulator_quantum_task.py

@@ -10,6 +10,7 @@
 # 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.
+
 import math
 
 import pytest

test/unit_tests/braket/annealing/test_problem.py

@@ -10,6 +10,7 @@
 # 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.
+
 import braket.ir.annealing as ir
 from braket.annealing.problem import Problem, ProblemType
 

test/unit_tests/braket/devices/test_local_simulator.py

@@ -10,8 +10,8 @@
 # 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.
-import json
-from typing import Optional
+
+from typing import Any, Dict, Optional
 from unittest.mock import Mock
 
 import braket.ir as ir
@@ -22,31 +22,32 @@
 from braket.devices.braket_simulator import BraketSimulator
 from braket.tasks import AnnealingQuantumTaskResult, GateModelQuantumTaskResult
 
-GATE_MODEL_RESULT = 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]],
-        "TaskMetadata": {"Id": "UUID_blah_1", "Status": "COMPLETED"},
-    }
-)
-
-ANNEALING_RESULT = json.dumps(
-    {
-        "Solutions": [[-1, -1, -1, -1], [1, -1, 1, 1], [1, -1, -1, 1]],
-        "VariableCount": 4,
-        "Values": [0.0, 1.0, 2.0],
-        "SolutionCounts": None,
-        "ProblemType": "ising",
-        "Foo": {"Bar": "Baz"},
-        "TaskMetadata": {"Id": "UUID_blah_1", "Status": "COMPLETED", "Shots": 5,},
-    }
-)
+GATE_MODEL_RESULT = {
+    "StateVector": {
+        "00": complex(0.2, 0.2),
+        "01": complex(0.3, 0.1),
+        "10": complex(0.1, 0.3),
+        "11": complex(0.2, 0.2),
+    },
+    "Measurements": [[0, 0], [0, 1], [0, 1], [0, 1]],
+    "TaskMetadata": {"Id": "UUID_blah_1", "Status": "COMPLETED"},
+}
+
+ANNEALING_RESULT = {
+    "Solutions": [[-1, -1, -1, -1], [1, -1, 1, 1], [1, -1, -1, 1]],
+    "VariableCount": 4,
+    "Values": [0.0, 1.0, 2.0],
+    "SolutionCounts": None,
+    "ProblemType": "ising",
+    "Foo": {"Bar": "Baz"},
+    "TaskMetadata": {"Id": "UUID_blah_1", "Status": "COMPLETED", "Shots": 5,},
+}
 
 
 class DummyCircuitSimulator(BraketSimulator):
     def run(
         self, program: ir.jaqcd.Program, qubits: int, shots: Optional[int], *args, **kwargs
-    ) -> str:
+    ) -> Dict[str, Any]:
         self._shots = shots
         self._qubits = qubits
         return GATE_MODEL_RESULT
@@ -59,7 +60,7 @@ def assert_qubits(self, qubits):
 
 
 class DummyAnnealingSimulator(BraketSimulator):
-    def run(self, problem: ir.annealing.Problem, *args, **kwargs) -> str:
+    def run(self, problem: ir.annealing.Problem, *args, **kwargs) -> Dict[str, Any]:
         return ANNEALING_RESULT
 
 
@@ -71,7 +72,7 @@ def run(self, problem: ir.annealing.Problem, *args, **kwargs) -> str:
 def test_load_from_entry_point():
     sim = LocalSimulator("dummy")
     task = sim.run(Circuit().h(0).cnot(0, 1), 10)
-    assert task.result == GateModelQuantumTaskResult.from_string(GATE_MODEL_RESULT)
+    assert task.result == GateModelQuantumTaskResult.from_dict(GATE_MODEL_RESULT)
 
 
 def test_run_gate_model():
@@ -80,13 +81,13 @@ def test_run_gate_model():
     task = sim.run(Circuit().h(0).cnot(0, 1), 10)
     dummy.assert_shots(10)
     dummy.assert_qubits(2)
-    assert task.result == GateModelQuantumTaskResult.from_string(GATE_MODEL_RESULT)
+    assert task.result == GateModelQuantumTaskResult.from_dict(GATE_MODEL_RESULT)
 
 
 def test_run_annealing():
     sim = LocalSimulator(DummyAnnealingSimulator())
     task = sim.run(Problem(ProblemType.ISING))
-    assert task.result == AnnealingQuantumTaskResult.from_string(ANNEALING_RESULT)
+    assert task.result == AnnealingQuantumTaskResult.from_dict(ANNEALING_RESULT)
 
 
 def test_registered_backends():