Add local simulator support (#57)

* Add BraketSimulator class that will eventually be moved to the simulator repo. This is the abstract class that all local simulators should implement.
* Added __init__.py for braket.annealing module
* Also fixed a couple of malformatted docstrings

src/braket/annealing/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+# Execute initialization code in circuit module
+from braket.annealing.problem import Problem, ProblemType  # noqa: F401

src/braket/aws/aws_quantum_task.py

@@ -154,11 +154,12 @@ def cancel(self) -> None:
     def metadata(self, use_cached_value: bool = False) -> Dict[str, Any]:
         """
         Get task metadata defined in Amazon Braket.
+
         Args:
             use_cached_value (bool, optional): If `True`, uses the value most recently retrieved
-            from the Amazon Braket `GetQuantumTask` operation. If `False`, calls the
-            `GetQuantumTask` operation  to retrieve metadata, which also updates the cached
-            value. Default = False.
+                from the Amazon Braket `GetQuantumTask` operation. If `False`, calls the
+                `GetQuantumTask` operation  to retrieve metadata, which also updates the cached
+                value. Default = False.
         Returns:
             Dict[str, Any]: The response from the Amazon Braket `GetQuantumTask` operation.
             If `use_cached_value` is `True`, Amazon Braket is not called and the most recently
@@ -171,11 +172,12 @@ def metadata(self, use_cached_value: bool = False) -> Dict[str, Any]:
     def state(self, use_cached_value: bool = False) -> str:
         """
         The state of the quantum task.
+
         Args:
             use_cached_value (bool, optional): If `True`, uses the value most recently retrieved
-            from the Amazon Braket `GetQuantumTask` operation. If `False`, calls the
-            `GetQuantumTask` operation to retrieve metadata, which also updates the cached
-            value. Default = False.
+                from the Amazon Braket `GetQuantumTask` operation. If `False`, calls the
+                `GetQuantumTask` operation to retrieve metadata, which also updates the cached
+                value. Default = False.
         Returns:
             str: The value of `status` in `metadata()`. This is the value of the `status` key
             in the Amazon Braket `GetQuantumTask` operation. If `use_cached_value` is `True`,
@@ -225,14 +227,17 @@ async def _create_future(self) -> asyncio.Task:
         """
         return asyncio.create_task(self._wait_for_completion())
 
-    async def _wait_for_completion(self) -> GateModelQuantumTaskResult:
+    async def _wait_for_completion(
+        self,
+    ) -> Union[GateModelQuantumTaskResult, AnnealingQuantumTaskResult]:
         """
         Waits for the quantum task to be completed, then returns the result from the S3 bucket.
         Returns:
-            GateModelQuantumTaskResult: If the task is in the `AwsQuantumTask.RESULTS_READY_STATES`
-                state within the specified time limit, the result from the S3 bucket is loaded and
-                returned. `None` is returned if a timeout occurs or task state is in
-                `AwsQuantumTask.TERMINAL_STATES` but not `AwsQuantumTask.RESULTS_READY_STATES`.
+            Union[GateModelQuantumTaskResult, AnnealingQuantumTaskResult]: If the task is in
+                the `AwsQuantumTask.RESULTS_READY_STATES` state within the specified time limit,
+                the result from the S3 bucket is loaded and returned. `None` is returned if a
+                timeout occurs or task state is in `AwsQuantumTask.TERMINAL_STATES` but not
+                `AwsQuantumTask.RESULTS_READY_STATES`.
         Note:
             Timeout and sleep intervals are defined in the constructor fields
                 `poll_timeout_seconds` and `poll_interval_seconds` respectively.

src/braket/devices/__init__.py

@@ -12,3 +12,4 @@
 # language governing permissions and limitations under the License.
 
 from braket.devices.device import Device  # noqa: F401
+from braket.devices.local_simulator import LocalSimulator  # noqa: F401

src/braket/devices/braket_simulator.py

@@ -0,0 +1,59 @@
+# 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 abc import ABC, abstractmethod
+from typing import Union
+
+from braket.ir.annealing import Problem
+from braket.ir.jaqcd import Program
+
+
+class BraketSimulator(ABC):
+    """ An abstract simulator that locally runs a quantum task.
+
+    The task can be either a circuit-based program or an annealing task,
+    specified by the given IR.
+
+    For users creating their own simulator: to register a simulator so the
+    Braket SDK recognizes its name, the name and class must added as an
+    entry point for "braket.simulators". This is done by adding an entry to
+    entry_points in the simulator package's setup.py:
+
+    >>> entry_points = {
+    >>>     "braket.simulators": [
+    >>>         "backend_name = <backend_class>"
+    >>>     ]
+    >>> }
+    """
+
+    # TODO: Move this class to the local simulator repo and take a dependency on it
+    # As such, this will not depend on any SDK classes.
+
+    # TODO: Update to use new simulate() method
+
+    @abstractmethod
+    def run(self, ir: Union[Program, Problem], *args, **kwargs) -> str:
+        """ Run the task specified by the given IR.
+
+        Extra arguments will contain any additional information necessary to run the task,
+        such as number of qubits.
+
+        Args:
+            ir (Union[Program, Problem]): The IR representation of the program
+
+        Returns:
+            str: A JSON string 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.
+        """
+        raise NotImplementedError()

src/braket/devices/device.py

@@ -49,7 +49,7 @@ def run(
 
         Args:
             task_specification (Union[Circuit, Problem]):  Specification of a task
-            to run on device.
+                to run on device.
 
             location: The location to save the task's results
 

src/braket/devices/local_simulator.py

@@ -0,0 +1,119 @@
+# 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 functools import singledispatch
+from typing import Set, Union
+
+import pkg_resources
+
+from braket.annealing.problem import Problem
+from braket.circuits import Circuit
+from braket.devices.braket_simulator import BraketSimulator
+from braket.devices.device import Device
+from braket.tasks import AnnealingQuantumTaskResult, GateModelQuantumTaskResult
+from braket.tasks.local_quantum_task import LocalQuantumTask
+
+_simulator_devices = {
+    entry.name: entry for entry in pkg_resources.iter_entry_points("braket.simulators")
+}
+
+
+class LocalSimulator(Device):
+    """ A simulator meant to run directly on the user's machine.
+
+    This class wraps a BraketSimulator object so that it can be run and returns
+    results using constructs from the SDK rather than Braket IR.
+    """
+
+    def __init__(self, backend: Union[str, BraketSimulator]):
+        """
+        Args:
+            backend (Union[str, BraketSimulator]): The name of the simulator backend or
+                the actual simulator instance to use for simulation
+        """
+        delegate = _get_simulator(backend)
+        super().__init__(
+            name=delegate.__class__.__name__,
+            status="AVAILABLE",
+            status_reason="Local simulator loaded successfully",
+        )
+        self._delegate = delegate
+
+    def run(
+        self, task_specification: Union[Circuit, Problem], *args, **kwargs,
+    ) -> LocalQuantumTask:
+        """ Runs the given task with the wrapped local simulator.
+
+        Args:
+            task_specification (Union[Circuit, Problem]):
+            *args: Positional args to pass to the IR simulator
+            **kwargs: Keyword arguments to pass to the IR simulator
+
+        Returns:
+            LocalQuantumTask: A LocalQuantumTask object containing the results
+            of the simulation
+
+        Note:
+            If running a circuit, the number of qubits will be passed
+            to the backend as the argument after the circuit itself.
+        """
+        result = _run_internal(task_specification, self._delegate, *args, **kwargs)
+        return LocalQuantumTask(result)
+
+    @classmethod
+    def registered_backends(cls) -> Set[str]:
+        """ Gets the backends that have been registered as entry points
+
+        Returns:
+            Set[str]: The names of the available backends that can be passed
+            into LocalSimulator's constructor
+        """
+        return set(_simulator_devices.keys())
+
+
+@singledispatch
+def _get_simulator(simulator):
+    raise TypeError("Simulator must either be a string or a BraketSimulator instance")
+
+
+@_get_simulator.register
+def _(backend_name: str):
+    if backend_name in _simulator_devices:
+        device_class = _simulator_devices[backend_name].load()
+        return device_class()
+    else:
+        raise ValueError(f"Only the following devices are available {_simulator_devices.keys()}")
+
+
+@_get_simulator.register
+def _(backend_impl: BraketSimulator):
+    return backend_impl
+
+
+@singledispatch
+def _run_internal(task_specification, simulator: BraketSimulator, *args, **kwargs):
+    raise NotImplementedError("Unsupported task type")
+
+
+@_run_internal.register
+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)
+
+
+@_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)

src/braket/tasks/annealing_quantum_task_result.py

@@ -28,9 +28,9 @@ class AnnealingQuantumTaskResult:
 
     Args:
         record_array (numpy.recarray): numpy array with keys 'solution' (numpy.ndarray)
-        where row is solution, column is value of the variable, 'solution_count' (numpy.ndarray)
-        the number of times the solutions occurred, and 'value' (numpy.ndarray) the
-        output or energy of the solutions.
+            where row is solution, column is value of the variable, 'solution_count' (numpy.ndarray)
+            the number of times the solutions occurred, and 'value' (numpy.ndarray) the
+            output or energy of the solutions.
         variable_count (int): the number of variables
         problem_type (str): the type of problem ('ising' or 'qubo')
         task_metadata (Dict[str, Any]): Dictionary of task metadata.

src/braket/tasks/gate_model_quantum_task_result.py

@@ -140,6 +140,10 @@ def from_string(result: str) -> GateModelQuantumTaskResult:
 
         Returns:
             GateModelQuantumTaskResult: A GateModelQuantumTaskResult based on a string
+
+        Raises:
+            ValueError: If neither "Measurements" nor "MeasurementProbabilities" is a key
+            in the result dict
         """
         json_obj = json.loads(result)
         task_metadata = json_obj["TaskMetadata"]
@@ -166,6 +170,10 @@ def from_string(result: str) -> GateModelQuantumTaskResult:
             measurements_copied_from_device = False
             m_counts_copied_from_device = False
             m_probabilities_copied_from_device = True
+        else:
+            raise ValueError(
+                'One of "Measurements" or "MeasurementProbabilities" must be in the results dict'
+            )
         return GateModelQuantumTaskResult(
             state_vector=state_vector,
             task_metadata=task_metadata,

src/braket/tasks/local_quantum_task.py

@@ -0,0 +1,46 @@
+# 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.
+import asyncio
+import uuid
+from typing import Union
+
+from braket.tasks import AnnealingQuantumTaskResult, GateModelQuantumTaskResult, QuantumTask
+
+
+class LocalQuantumTask(QuantumTask):
+    """ A task containing the results of a local simulation.
+
+    Since this class is instantiated with the results, cancel() and run_async() are unsupported.
+    """
+
+    def __init__(self, result: Union[GateModelQuantumTaskResult, AnnealingQuantumTaskResult]):
+        self._id = uuid.uuid4()
+        self._result = result
+
+    @property
+    def id(self) -> str:
+        return str(self._id)
+
+    def cancel(self) -> None:
+        raise NotImplementedError("Cannot cancel completed local task")
+
+    def state(self) -> str:
+        return "COMPLETED"
+
+    @property
+    def result(self) -> Union[GateModelQuantumTaskResult, AnnealingQuantumTaskResult]:
+        return self._result
+
+    def async_result(self) -> asyncio.Task:
+        # TODO: Allow for asynchronous simulation
+        raise NotImplementedError("Asynchronous local simulation unsupported")

src/braket/tasks/quantum_task.py

@@ -52,7 +52,7 @@ def metadata(self, use_cached_value: bool = False) -> Dict[str, Any]:
 
         Args:
             use_cached_value (bool, optional): If True, uses the value retrieved from the previous
-            request.
+                request.
 
         Returns:
             Dict[str, Any]: The metadata regarding the task. If `use_cached_value` is True,