feature: Add simulator module with mcm simulator implementation (#17)

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Changelog
 
-## v0.1.0 (2024-05-09)
+## v0.1.0 (2024-05-22)
 
 This is the initial pre-release version of AutoQASM.
 

README.md

@@ -96,13 +96,13 @@ AutoQASM can support subroutines and complex control flow. You can use the Pytho
 and quantum runtime side-by-side. There are rough edges at the moment, but we're actively smoothing
 them out!
 
-The Amazon Braket local simulator supports AutoQASM programs as input.
+AutoQASM includes a simulator which can be accessed using the Amazon Braket local simulator interface.
 Let's simulate the `conditional_multi_bell_states` program:
 
 ```
-from braket.devices.local_simulator import LocalSimulator
+from braket.devices import LocalSimulator
 
-device = LocalSimulator()
+device = LocalSimulator("autoqasm")
 task = device.run(conditional_multi_bell_states, shots=100)
 result = task.result()
 ```

examples/1_Getting_started_with_AutoQASM.ipynb

@@ -21,7 +21,7 @@
     "import matplotlib.pyplot as plt\n",
     "\n",
     "# AWS imports: Import Braket SDK modules\n",
-    "from braket.devices.local_simulator import LocalSimulator\n",
+    "from braket.devices import LocalSimulator\n",
     "\n",
     "# AutoQASM imports\n",
     "import autoqasm as aq\n",
@@ -108,7 +108,7 @@
     }
    ],
    "source": [
-    "device = LocalSimulator()\n",
+    "device = LocalSimulator(\"autoqasm\")\n",
     "result = device.run(bell_state, shots=100).result()\n",
     "counts = Counter(result.measurements[\"return_value\"])\n",
     "print(\"measurement counts: \", counts)"

examples/2_Expressing_classical_control_flow.ipynb

@@ -22,7 +22,7 @@
     "import matplotlib.pyplot as plt\n",
     "\n",
     "# AWS imports: Import Braket SDK modules\n",
-    "from braket.devices.local_simulator import LocalSimulator\n",
+    "from braket.devices import LocalSimulator\n",
     "\n",
     "# AutoQASM imports\n",
     "import autoqasm as aq\n",
@@ -167,7 +167,7 @@
     }
    ],
    "source": [
-    "device = LocalSimulator()\n",
+    "device = LocalSimulator(\"autoqasm\")\n",
     "result = device.run(conditioned_bell, shots=500).result()\n",
     "counts = Counter(result.measurements[\"return_value\"])\n",
     "print(\"measurement counts: \", counts)\n",

examples/3_1_Iterative_phase_estimation.ipynb

@@ -24,7 +24,7 @@
     "import matplotlib.pyplot as plt\n",
     "\n",
     "# AWS imports: Import Braket SDK modules\n",
-    "from braket.devices.local_simulator import LocalSimulator\n",
+    "from braket.devices import LocalSimulator\n",
     "\n",
     "# AutoQASM imports\n",
     "import autoqasm as aq\n",
@@ -158,7 +158,7 @@
     }
    ],
    "source": [
-    "device = LocalSimulator()\n",
+    "device = LocalSimulator(\"autoqasm\")\n",
     "result = device.run(ipe, shots=100).result()\n",
     "counts = Counter(result.measurements[\"b0\"])\n",
     "print(\"measurement counts: \", counts)\n",

examples/3_2_magic_state_distillation.ipynb

@@ -37,7 +37,7 @@
     "from collections import defaultdict\n",
     "\n",
     "# AWS imports: Import Braket SDK modules\n",
-    "from braket.devices.local_simulator import LocalSimulator\n",
+    "from braket.devices import LocalSimulator\n",
     "\n",
     "# AutoQASM imports\n",
     "import autoqasm as aq\n",
@@ -175,7 +175,7 @@
    ],
    "source": [
     "# Get measurement result\n",
-    "result = LocalSimulator().run(gate_teleportation, shots=100).result()\n",
+    "result = LocalSimulator(\"autoqasm\").run(gate_teleportation, shots=100).result()\n",
     "counts = Counter(result.measurements[\"return_value\"])\n",
     "print(\"measurement counts: \", counts)\n",
     "\n",
@@ -315,7 +315,7 @@
    ],
    "source": [
     "n_shots = 1000\n",
-    "result = LocalSimulator().run(distillation, shots=n_shots).result()\n",
+    "result = LocalSimulator(\"autoqasm\").run(distillation, shots=n_shots).result()\n",
     "counts = Counter(result.measurements[\"c\"])\n",
     "print(\"measurement counts: \", counts)"
    ]
@@ -435,7 +435,7 @@
     }
    ],
    "source": [
-    "result = LocalSimulator().run(distillation_rus, shots=20).result()\n",
+    "result = LocalSimulator(\"autoqasm\").run(distillation_rus, shots=20).result()\n",
     "counts = Counter(result.measurements[\"c2\"])\n",
     "probs = {str(k): v / sum(counts.values()) for k, v in counts.items()}\n",
     "\n",

setup.py

@@ -24,12 +24,8 @@
     packages=find_namespace_packages(where="src", exclude=("test",)),
     package_dir={"": "src"},
     install_requires=[
-        # Pin the latest commit of mcm-sim branch of amazon-braket/amazon-braket-sdk-python.git
-        # and amazon-braket/amazon-braket-default-simulator-python.git to get the version of the
-        # simulator that supports the mcm=True argument for Monte Carlo simulation of mid-circuit
-        # measurement, which AutoQASM requires.
-        "amazon-braket-sdk @ git+https://github.com/amazon-braket/amazon-braket-sdk-python.git@ff73de68cf6ac2d0a921e8fe62693e5b9ae2e321#egg=amazon-braket-sdk",  # noqa E501
-        "amazon-braket-default-simulator @ git+https://github.com/amazon-braket/amazon-braket-default-simulator-python.git@ab068c860963c29842d7649c741f88da669597eb#egg=amazon-braket-default-simulator",  # noqa E501
+        "amazon-braket-sdk>=1.80.0",
+        "amazon-braket-default-simulator>=1.23.2",
         "oqpy~=0.3.5",
         "diastatic-malt",
         "numpy",
@@ -60,6 +56,11 @@
             "tox",
         ],
     },
+    entry_points={
+        "braket.simulators": [
+            "autoqasm = autoqasm.simulator.simulator:McmSimulator",
+        ]
+    },
     include_package_data=True,
     url="https://github.com/amazon-braket/autoqasm",
     author="Amazon Web Services",

src/autoqasm/program/program.py

@@ -121,28 +121,28 @@ def build(self, device: Device | str | None = None) -> Program:
     def to_ir(
         self,
         ir_type: IRType = IRType.OPENQASM,
-        allow_implicit_build: bool = False,
+        build_if_necessary: bool = True,
         serialization_properties: SerializationProperties = OpenQASMSerializationProperties(),
     ) -> str:
         """Serializes the program into an intermediate representation.
 
         Args:
             ir_type (IRType): The IRType to use for converting the program to its
                 IR representation. Defaults to IRType.OPENQASM.
-            allow_implicit_build (bool): Whether to allow the program to be implicitly
-                built as a side effect of calling this function. Defaults to False.
+            build_if_necessary (bool): Whether to allow the program to be implicitly
+                built as a side effect of calling this function. Defaults to True.
             serialization_properties (SerializationProperties): IR serialization configuration.
                 Default to OpenQASMSerializationProperties().
 
         Raises:
             ValueError: Raised if the supplied `ir_type` is not supported.
-            RuntimeError: Raised if `allow_implicit_build` is False, since a MainProgram object
+            RuntimeError: Raised if `build_if_necessary` is False, since a MainProgram object
                 has not yet been built.
 
         Returns:
             str: A representation of the program in the `ir_type` format.
         """
-        if not allow_implicit_build:
+        if not build_if_necessary:
             raise RuntimeError(
                 "The AutoQASM program cannot be serialized because it has not yet been built. "
                 "To serialize the program, first call build() to obtain a built Program object, "
@@ -227,15 +227,15 @@ def make_bound_program(self, param_values: dict[str, float], strict: bool = Fals
     def to_ir(
         self,
         ir_type: IRType = IRType.OPENQASM,
-        allow_implicit_build: bool = True,
+        build_if_necessary: bool = True,
         serialization_properties: SerializationProperties = OpenQASMSerializationProperties(),
     ) -> str:
         """Serializes the program into an intermediate representation.
 
         Args:
             ir_type (IRType): The IRType to use for converting the program to its
                 IR representation. Defaults to IRType.OPENQASM.
-            allow_implicit_build (bool): Whether to allow the program to be implicitly
+            build_if_necessary (bool): Whether to allow the program to be implicitly
                 built as a side effect of calling this function. Defaults to True.
                 This parameter is ignored for the Program class, since the program has
                 already been built.

src/autoqasm/simulator/__init__.py

@@ -0,0 +1,37 @@
+# Copyright 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.
+
+"""
+This module contains a local simulator implementation which can be used to
+simulate AutoQASM programs at the full OpenQASM 3.0 scope of functionality,
+including programs which contain mid-circuit measurement and classical control flow.
+
+The recommended usage of this simulator is by instantiating the Braket LocalSimulator
+object with the "autoqasm" backend:
+
+.. code-block:: python
+
+    import autoqasm as aq
+    from braket.devices import LocalSimulator
+
+    @aq.main
+    def bell_state():
+        h(0)
+        cnot(0, 1)
+        return measure([0, 1])
+
+    device = LocalSimulator("autoqasm")
+    result = device.run(bell_state, shots=100).result()
+"""
+
+from .simulator import McmSimulator  # noqa: F401

src/autoqasm/simulator/conversion.py

@@ -0,0 +1,47 @@
+# Copyright 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 Any, Union
+
+import numpy as np
+from braket.default_simulator.openqasm._helpers.casting import convert_bool_array_to_string
+from braket.default_simulator.openqasm.parser.openqasm_ast import (
+    ArrayLiteral,
+    BitstringLiteral,
+    BooleanLiteral,
+    FloatLiteral,
+    IntegerLiteral,
+)
+
+LiteralType = Union[BooleanLiteral, IntegerLiteral, FloatLiteral, ArrayLiteral, BitstringLiteral]
+
+
+@singledispatch
+def convert_to_output(value: LiteralType) -> Any:
+    raise NotImplementedError(f"converting {value} to output")
+
+
+@convert_to_output.register(IntegerLiteral)
+@convert_to_output.register(FloatLiteral)
+@convert_to_output.register(BooleanLiteral)
+@convert_to_output.register(BitstringLiteral)
+def _(value):
+    return value.value
+
+
+@convert_to_output.register
+def _(value: ArrayLiteral):
+    if isinstance(value.values[0], BooleanLiteral):
+        return convert_bool_array_to_string(value)
+    return np.array([convert_to_output(x) for x in value.values])

src/autoqasm/simulator/linalg_utils.py

@@ -0,0 +1,56 @@
+# Copyright 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 itertools
+from collections.abc import Iterable
+
+import numpy as np
+
+
+def measurement_sample(prob: Iterable[float], target_count: int) -> tuple[int]:
+    """Draws a random measurement sample.
+
+    Args:
+        prob (Iterable[float]): Probabilities of each measurement outcome.
+            Possible measurement outcomes range from 0 to 2**target_count-1,
+            meaning that len(prob) must be equal to 2**target_count.
+        target_count (int): Number of bits in the resulting measurement.
+
+    Returns:
+        tuple[int]: Measurement outcomes 0 or 1 for each of the target_count bits.
+    """
+    basis_states = np.array(list(itertools.product([0, 1], repeat=target_count)))
+    outcome_idx = np.random.choice(list(range(2**target_count)), p=prob)
+    return tuple(basis_states[outcome_idx])
+
+
+def measurement_collapse_sv(
+    state_vector: np.ndarray, targets: Iterable[int], outcome: np.ndarray
+) -> np.ndarray:
+    """Collapses the state vector according to the given measurement outcome.
+
+    Args:
+        state_vector (np.ndarray): The state vector prior to measurement.
+        targets (Iterable[int]): The qubit indices that were measured.
+        outcome (np.ndarray): Array of measurement outcomes 0 or 1.
+
+    Returns:
+        np.ndarray: The collapsed state vector after measurement.
+    """
+    qubit_count = int(np.log2(state_vector.size))
+    state_tensor = state_vector.reshape([2] * qubit_count)
+    for qubit, measurement in zip(targets, outcome):
+        state_tensor[(slice(None),) * qubit + (int(not measurement),)] = 0
+
+    state_tensor /= np.linalg.norm(state_tensor)
+    return state_tensor.flatten()