feature: Approximate State Preparation using Sweeping

README.md

@@ -26,7 +26,7 @@ Running notebooks locally requires additional dependencies located in [notebooks
 | Quantum Walk | [Quantum_Walk.ipynb](notebooks/textbook/Quantum_Walk.ipynb) | [Childs2002](https://arxiv.org/abs/quant-ph/0209131) |
 |Shor's| [Shors_Algorithm.ipynb](notebooks/textbook/Shors_Algorithm.ipynb) | [Shor1998](https://arxiv.org/abs/quant-ph/9508027) |
 | Simon's | [Simons_Algorithm.ipynb](notebooks/textbook/Simons_Algorithm.ipynb) | [Simon1997](https://epubs.siam.org/doi/10.1137/S0097539796298637) |
-
+| Sweeping | [Sweeping.ipynb](notebooks/textbook/Sweeping.ipynb) | [Rudolph2022](https://arxiv.org/abs/2209.00595) |
 
 | Advanced algorithms | Notebook | References |
 | ----- | ----- | ----- |

src/braket/experimental/algorithms/sweeping/__init__.py

@@ -0,0 +1,24 @@
+# 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.
+
+__all__ = [
+    "generate_brickwall_ansatz",
+    "generate_staircase_ansatz",
+    "sweep_state_approximation",
+]
+
+from braket.experimental.algorithms.sweeping.ansatzes import (
+    generate_brickwall_ansatz,
+    generate_staircase_ansatz,
+)
+from braket.experimental.algorithms.sweeping.sweeping import sweep_state_approximation

src/braket/experimental/algorithms/sweeping/ansatzes.py

@@ -0,0 +1,54 @@
+import numpy as np
+
+from braket.experimental.algorithms.sweeping.typing import UNITARY_LAYER
+
+
+def generate_staircase_ansatz(num_qubits: int, num_layers: int) -> list[UNITARY_LAYER]:
+    """Generate a staircase ansatz with the specified number of qubits and layers.
+
+    Args:
+        num_qubits (int): The number of qubits in the ansatz.
+        num_layers (int): The number of layers in the ansatz.
+
+    Returns:
+        list[UNITARY_LAYER]: The generated staircase ansatz.
+    """
+    unitary_layers: list[UNITARY_LAYER] = []
+
+    for i in range(num_layers):
+        unitary_layers.append([])
+
+        for j in range(num_qubits - 1):
+            unitary_layers[i].append(([j, j + 1], np.eye(4, dtype=np.complex128)))
+
+    return unitary_layers
+
+
+def generate_brickwall_ansatz(num_qubits: int, num_layers: int) -> list[UNITARY_LAYER]:
+    """Generate a brickwall ansatz with the specified number of qubits and layers.
+
+    Args:
+        num_qubits (int): The number of qubits in the ansatz.
+        num_layers (int): The number of layers in the ansatz.
+
+    Returns:
+        list[UNITARY_LAYER]: The generated brickwall ansatz.
+    """
+    unitary_layers: list[UNITARY_LAYER] = []
+
+    if num_qubits % 2 == 0:
+        start_1 = 1
+        start_2 = 0
+    else:
+        start_1 = 0
+        start_2 = 1
+
+    for i in range(num_layers):
+        unitary_layers.append([])
+
+        for j in range(start_1, num_qubits - 1, 2):
+            unitary_layers[i].append(([j, j + 1], np.eye(4, dtype=np.complex128)))
+        for j in range(start_2, num_qubits - 1, 2):
+            unitary_layers[i].append(([j, j + 1], np.eye(4, dtype=np.complex128)))
+
+    return unitary_layers

src/braket/experimental/algorithms/sweeping/sweeping.md

@@ -0,0 +1,10 @@
+Tensor networks are a natural way to interact with quantum circuits and develop algorithms for them. A major application of tensor networks is in performing circuit optimization,
+either via MPS/MPO IRs, or via sweeping. We focus on using tensor network sweeping as an algorithm to optimize arbitrary ansatzes to approximate a target statevector or unitary
+in a smooth, consistently improving manner. The approach is a much better alternative for such cases compared to gradient-based and ML approaches which are prone to local minimas
+and barren plateaus. For a better understanding of the technique see [Rudolph2022](https://arxiv.org/abs/2209.00595).
+
+<!--
+[metadata-name]: Sweeping
+[metadata-tags]: Textbook
+[metadata-url]: https://github.com/amazon-braket/amazon-braket-algorithm-library/tree/main/src/braket/experimental/algorithms/sweeping
+-->

src/braket/experimental/algorithms/sweeping/sweeping.py

@@ -0,0 +1,244 @@
+import numpy as np
+import quimb.tensor as qtn
+from numpy.typing import NDArray
+from tqdm import tqdm
+
+from braket.circuits import Circuit
+from braket.devices import LocalSimulator
+from braket.experimental.algorithms.sweeping.typing import UNITARY_LAYER
+
+
+class StateSweepApproximator:
+    """Initialize the StateSweepApproximator class.
+
+    This algorithm performs sweeps over unitary layers representing
+    an ansatz to be trained to match the target state to improve
+    its approximation. Due to the difficulty in converging to good
+    fidelity with pure classical ML approaches, and the time it takes
+    to train them, sweeping is a much better alternative, running faster
+    and giving smooth improvement.
+
+    This approach is inspired by Rudolph et al. [1] which was used to do
+    the same task under Oall approach.
+
+    [1] https://www.nature.com/articles/s41467-023-43908-6
+    """
+
+    @staticmethod
+    def get_tensor_network_from_unitary_layers(
+        num_qubits: int, unitary_layers: list[UNITARY_LAYER]
+    ) -> qtn.TensorNetwork:
+        """Create `qtn.TensorNetwork` from unitary layers.
+
+        Args:
+            num_qubits (int): The number of qubits used by target state.
+            unitary_layers (list[UNITARY_LAYER]): The unitary layers.
+
+        Returns:
+            tensor_network (qtn.TensorNetwork): The tensor network.
+        """
+        circuit = qtn.Circuit(N=num_qubits)
+        gate_tracker: list[str] = []
+
+        for i, layer in enumerate(unitary_layers):
+            for j, (qubits, unitary) in enumerate(layer):
+                circuit.apply_gate_raw(
+                    unitary.reshape(2 * len(qubits) * (2,)),
+                    where=qubits,
+                    contract=False,
+                )
+                gate_tracker.append(f"{i}_{j}")
+
+        tensor_network = qtn.TensorNetwork(circuit.psi)
+
+        # We do not want to include the qubits, so we will
+        # explicitly control the iteration index
+        gate_index = 0
+
+        for gate in tensor_network:
+            # We only update the gates, not the qubits
+            if "PSI0" in gate.tags:
+                continue
+
+            # Remove existing tags from the gate
+            gate.drop_tags(tags=gate.tags)
+
+            # Marshal the gate with the gate tracker
+            # This is needed to ensure the gates are properly tagged
+            # for updating the unitary layers
+            gate.add_tag(gate_tracker[gate_index])
+            gate_index += 1
+
+        return tensor_network
+
+    @staticmethod
+    def sweep_unitary_layers(
+        target_mps: qtn.MatrixProductState,
+        tensor_network: qtn.TensorNetwork,
+        unitary_layers: list[UNITARY_LAYER],
+    ) -> list[UNITARY_LAYER]:
+        """Sweep the unitary layers to improve the fidelity between
+        tensor network created by the unitary layers and the target
+        state.
+
+        Args:
+            target_mps (qtn.MatrixProductState): The target state represented as a MPS.
+            tensor_network (qtn.TensorNetwork): The tensor network created by the unitary layers.
+            unitary_layers (list[UNITARY_LAYER]): The unitary layers.
+
+        Returns:
+            unitary_layers (list[UNITARY_LAYER]): The unitary layers which have been updated inplace.
+        """
+        target_mps_adjoint = target_mps.conj()
+
+        current_tn = qtn.MatrixProductState.from_dense(
+            tensor_network.to_dense(tensor_network.outer_inds())
+        )
+
+        for gate in reversed(tensor_network.tensors):
+            num_qubits = int(len(gate.inds) / 2)
+
+            if "PSI0" in gate.tags:
+                continue
+
+            left_inds = gate.inds[:num_qubits]
+            right_inds = gate.inds[num_qubits:]
+
+            # To avoid unnecessary contraction, we "move"
+            # through the tensor network by applying the adjoint
+            # of the tensor to the tensor network, and applying
+            # the updated tensor to MPS
+            current_tn = current_tn @ gate.conj()
+            environment_tensor: qtn.TensorNetwork = target_mps_adjoint @ current_tn
+
+            u, _, vh = np.linalg.svd(environment_tensor.to_dense((left_inds), (right_inds)))
+            u_new = np.dot(u, vh)
+
+            # To avoid unnecessary contraction, we "move"
+            # through the tensor network by applying the adjoint
+            # of the tensor to the tensor network, and applying
+            # the updated tensor to MPS
+            new_tensor = qtn.Tensor(
+                u_new.reshape(2 * num_qubits * (2,)).conj(),
+                inds=gate.inds,
+                tags=gate.tags,
+            )
+
+            target_mps_adjoint = target_mps_adjoint @ new_tensor
+
+            gate_tag = list(gate.tags)[0]
+            layer_index, block_index = gate_tag.split("_")
+
+            unitary_layers[int(layer_index)][int(block_index)] = (
+                unitary_layers[int(layer_index)][int(block_index)][0],
+                u_new.conj(),
+            )
+
+        return unitary_layers
+
+    @staticmethod
+    def circuit_from_unitary_layers(
+        num_qubits: int, unitary_layers: list[UNITARY_LAYER]
+    ) -> Circuit:
+        """Create a `Circuit` instance from
+        the unitary layers.
+
+        Args:
+            num_qubits (int): The number of qubits used by target state.
+            unitary_layers (list[UNITARY_LAYER]): The unitary layers.
+
+        Returns:
+            circuit (Circuit): The qiskit circuit.
+        """
+        circuit = Circuit()
+
+        for layer in unitary_layers:
+            for qubits, unitary_matrix in layer:
+                circuit.unitary(
+                    matrix=unitary_matrix,
+                    targets=qubits,
+                )
+
+        return circuit
+
+    def __call__(
+        self,
+        target_state: NDArray[np.complex128],
+        unitary_layers: list[UNITARY_LAYER],
+        num_sweeps: int,
+        log: bool = False,
+    ) -> Circuit:
+        """Approximate a state via sweeping.
+
+        Args:
+            target_state (NDArray[np.complex128]): The state we want to approximate.
+            unitary_layers (list[UNITARY_LAYER]): The initial unitary layers.
+            num_sweeps (int): The number of times to sweep the unitary layers.
+            log (bool): Whether to print logs of fidelity or not.
+
+        Returns:
+            circuit (Circuit): The circuit.
+        """
+        num_qubits = int(np.log2(target_state.shape[0]))
+        target_mps = qtn.MatrixProductState.from_dense(target_state)
+
+        for i in tqdm(range(num_sweeps)):
+            circuit_tensor_network = self.get_tensor_network_from_unitary_layers(
+                num_qubits, unitary_layers
+            )
+
+            if log and i % 20 == 0:
+                fidelity = (
+                    np.abs(
+                        np.vdot(
+                            target_state,
+                            circuit_tensor_network.to_dense(
+                                circuit_tensor_network.outer_inds()
+                            ).reshape(target_state.shape),
+                        )
+                    )
+                    ** 2
+                )
+                print(f"Fidelity: {fidelity}")
+
+            unitary_layers = self.sweep_unitary_layers(
+                target_mps, circuit_tensor_network, unitary_layers
+            )
+
+        circuit = self.circuit_from_unitary_layers(num_qubits, unitary_layers)
+
+        if log:
+            result = LocalSimulator("braket_sv").run(circuit.state_vector(), shots=0).result()
+            fidelity = (
+                np.abs(
+                    np.vdot(
+                        result.values[0],
+                        target_state,
+                    )
+                )
+                ** 2
+            )
+            print(f"Final Fidelity: {fidelity}")
+
+        return circuit
+
+
+def sweep_state_approximation(
+    target_state: NDArray[np.complex128],
+    unitary_layers: list[UNITARY_LAYER],
+    num_sweeps: int,
+    log: bool = False,
+) -> Circuit:
+    """Approximate a state via sweeping.
+
+    Args:
+        target_state (NDArray[np.complex128]): The state we want to approximate.
+        unitary_layers (list[UNITARY_LAYER]): The initial unitary layers.
+        num_sweeps (int): The number of times to sweep the unitary layers.
+        log (bool): Whether to print logs of fidelity or not.
+
+    Returns:
+        circuit (Circuit): The circuit.
+    """
+    approximator = StateSweepApproximator()
+    return approximator(target_state, unitary_layers, num_sweeps, log)

src/braket/experimental/algorithms/sweeping/typing.py

@@ -0,0 +1,7 @@
+from typing import TypeAlias
+
+import numpy as np
+from numpy.typing import NDArray
+
+UNITARY_BLOCK: TypeAlias = tuple[list[int], NDArray[np.complex128]]
+UNITARY_LAYER: TypeAlias = list[UNITARY_BLOCK]

test/unit_tests/braket/experimental/algorithms/sweeping/test_sweeping.py

@@ -0,0 +1,28 @@
+import numpy as np
+import pytest
+
+from braket.devices import LocalSimulator
+from braket.experimental.algorithms.sweeping import (
+    generate_staircase_ansatz,
+    sweep_state_approximation,
+)
+
+
+@pytest.mark.parametrize("num_qubits", [2, 3, 4, 5])
+def compile_with_state_sweep_pass(num_qubits: int) -> None:
+    state = np.random.uniform(-1, 1, 2**num_qubits) + 1j * np.random.uniform(-1, 1, 2**num_qubits)
+    state /= np.linalg.norm(state)
+
+    compiled_circuit = sweep_state_approximation(
+        target_state=state,
+        unitary_layers=generate_staircase_ansatz(
+            num_qubits=num_qubits, num_layers=int((num_qubits**2) / 2)
+        ),
+        num_sweeps=100 * num_qubits,
+        log=False,
+    )
+
+    result = LocalSimulator("braket_sv").run(compiled_circuit.state_vector(), shots=0).result()
+    fidelity = np.abs(np.vdot(state, result.values[0]))
+
+    assert fidelity > 0.99