Ocean for Analog Hamiltonian Simulation

[christianbmadsen]

Modify https://github.com/aws/amazon-braket-ocean-plugin-python/ in a new branch such that you can embed a QUBO problem onto QuEra's Aquila machine via Braket.

Rewrite and demonstrate that existing examples can indeed run against Aquila with reasonable modifications (may need to for instance reduce the number of variables).

[JordanAWS]

If you want to work on this issue, please feel free to ping me to request AWS credits to cover the cost of running on the QuEra Aquila device :)

[Aaron-Robertson]

@JordanAWS I'm giving this one a shot! Could you send some AWS credits? I couldn't find your username on the discord to reach out (I'm A-Aron)

[JordanAWS]

Great! Can you send me an email at unitary-hack-braket-support@amazon.com with your AWS account ID? If you haven't yet created an account, please do so first :)

[Aaron-Robertson]

Apologies @JordanAWS, I got wrapped up in other issues for now! My goal is to come back to this one later today or tomorrow, but probably best to consider the "playing field" open for other hackers at the moment!

[JordanAWS]

@Aaron-Robertson Got it. It looks like you're currently the only one working on this issue, so it's all good if you'd like to come back to it later :)

[golanor]

I'd be happy to work on this with @Aaron-Robertson if he wants, or to work on it alone.

I want to make sure I understand the scope of the work here:

1. Create a new class BraketAquilaSampler in the plugin.
2. Modify the examples to work on Aquila.
3. Test on the device.

Is that correct @JordanAWS ? Any other things to take into consideration? (Other than asking for credits for testing the examples)

[Aaron-Robertson]

@golanor A collab sounds great to me! I'm not so sure we'll get by so easily, but if my hunch/brainstorming turns out correct, it will be a lot of fun!

First off, I think we'll need some metadata mapping, as is done in the BraketSolverMetadata.

Second, I'm looking in to a slightly different approach, though still just brainstorming for now. Ocean, as far as I can tell, does BQM based sampling under the hood, almost exclusively. Based on the resources on Aquila[1][2], it is similarly singularly focused, but on MIS problems. So, I'm wondering if what we need is really a topological embedding from QUBO to MIS, such that the ocean-based samplers directly (or as close as possible) interact with the embedded QUBO, which map to an Aquila MIS problem under the hood. Let me know what you think!

[Aaron-Robertson]

Alright, I need to look over everything here again with fresh eyes, but I stubbed out a possible approach on the embedding! Here's what I'm started on, though there's a ways to go:

import networkx as nx
import dimod
from dwave.system import EmbeddingComposite
import minorminer
from braket.ocean_plugin import BraketSampler
from braket.ahs.analog_hamiltonian_simulation import AnalogHamiltonianSimulation
from braket.ahs import Hamiltonian

class AquilaToOceanEmbeddingComposite(EmbeddingComposite):
    """
    This class extends the EmbeddingComposite class to include a method for
    solving the Maximum Independent Set (MIS) problem on an Aquila device.
    """
    def __init__(self, sampler):
        """
        Initialize the class with a sampler.
        
        Args:
            sampler: A sampler that can sample from binary quadratic models (BQMs).
        """
        super().__init__(sampler)

    def _bqm_to_mis(self, bqm):
        """
        Convert a Binary Quadratic Model (BQM) to a Maximum Independent Set (MIS).
        
        Args:
            bqm: A BQM to convert.
        
        Returns:
            A NetworkX graph representing the MIS.
        """
        # Construct a graph from the BQM
        graph = nx.Graph()
        for variable in bqm.variables:
            graph.add_node(variable)
        for edge, coeff in bqm.quadratic.items():
            graph.add_edge(*edge, weight=coeff)

        # Translate the MIS problem to a form that can be solved on an Aquila device
        aquila_mis = self._translate_to_aquila(graph)

        return aquila_mis


    def _mis_to_bqm(self, mis):
        """
        Convert a Maximum Independent Set (MIS) to a Binary Quadratic Model (BQM).
        
        Args:
            mis: A NetworkX graph representing the MIS.
        
        Returns:
            A BQM representing the MIS.
        """
        # Construct a BQM from the MIS
        linear = {v: 1 for v in mis}
        quadratic = {(i, j): 2 for i in mis for j in mis if i != j}
        bqm = dimod.BinaryQuadraticModel(linear, quadratic, 0, dimod.BINARY)
        return bqm

    def _translate_to_aquila(self, graph):
        """
        Translate a Maximum Independent Set (MIS) problem to a form that can be solved on an Aquila device.
        
        Args:
            graph: A NetworkX graph representing the MIS.
        
        Returns:
            A translated MIS problem.
        """
        # Create a new graph to represent the problem on the Aquila device
        aquila_graph = nx.Graph()
    
        # Map each node in the original graph to a Rydberg atom in the Aquila device
        for node in graph.nodes:
            aquila_graph.add_node(node)
    
        # Map each edge in the original graph to a connection between Rydberg atoms in the Aquila device
        for edge in graph.edges:
            aquila_graph.add_edge(*edge, weight=graph.edges[edge]['weight'])
    
        # Attempt to generate a unit disk graph from the aquila_graph
        try:
            unit_disk_graph = self._generate_unit_disk(aquila_graph)
        except Exception as e:
            print(f"An error occurred while generating the unit disk graph: {e}")
            raise
    
        # If the unit disk graph could not be generated, take another approach
        if unit_disk_graph is None:
            # TODO: Implement another approach for interfacing with the Aquila hardware
            pass
    
        # Return the translated graph
        return unit_disk_graph

    def _graph_to_ahs(self, graph):
        """
        Convert a graph to an Analog Hamiltonian Simulation (AHS) program.
        
        Args:
            graph: A NetworkX graph to convert.
        
        Returns:
            An AHS program representing the graph.
        """
        # Create a register of qubits
        register = list(range(len(graph.nodes)))

        # Create a Hamiltonian from the graph
        H = Hamiltonian(graph)

        # Create an AHS program from the Hamiltonian
        ahs_program = AnalogHamiltonianSimulation(register=register, hamiltonian=H)

        return ahs_program

    def _generate_unit_disk(self, graph):
      """
      Generate a unit disk graph from a given graph.
      
      Args:
          graph: A NetworkX graph to convert.
      
      Returns:
          A unit disk graph.
      """
      # Create a new graph to represent the unit disk graph
      unit_disk_graph = nx.Graph()
  
      # Iterate over the nodes in the input graph
      for node in graph.nodes:
      # Add the node to the unit disk graph
      unit_disk_graph.add_node(node)

      # Iterate over the edges in the input graph
      for edge in graph.edges:
          # Get the nodes at the ends of the edge
          node1, node2 = edge
      
          # Calculate the Euclidean distance between the nodes
          distance = np.sqrt((node1[0] - node2[0])**2 + (node1[1] - node2[1])**2)
      
          # If the distance is less than or equal to 1, add an edge between the nodes in the unit disk graph
          if distance <= 1:
              unit_disk_graph.add_edge(node1, node2)
          else:
              raise ValueError("The input graph is not a unit disk graph")

    # Return the unit disk graph
    return unit_disk_graph

    def sample(self, bqm):
        """
        Sample from a Binary Quadratic Model (BQM) using the underlying sampler.
        
        Args:
            bqm: A BQM to sample from.
        
        Returns:
            A sample set.
        """
        # TODO Need to actually solve the Quantum MIS with these samples. Still working out the details
        # TODO Probably also need to translate the samples back into a BQM? Something about using the samples to recreate a 
        BQM possibly using a variant of the above _mis_to_bqm() method.
        # Check that the input is a valid BQM
        if not isinstance(bqm, dimod.BinaryQuadraticModel):
            raise ValueError("Input must be a BinaryQuadraticModel")

        # Convert the BQM to an MIS problem
        mis = self._bqm_to_mis(bqm)

        ahs_program = self._graph_to_ahs(mis)
        
        # Probably unnecessary #
        # Find an embedding from the MIS problem to the quantum hardware
        target_graph = self.sampler.edgelist_to_adjacency(self.sampler.properties['couplers'])
        source_graph = nx.Graph(mis)
        embedding = minorminer.find_embedding(source_graph, target_graph)

        # Use the embedding to map the MIS problem to the quantum hardware
        embedded_bqm = dimod.embed_bqm(bqm, embedding, target_graph)
        # Probably unnecessary #

        # Try to sample from the MIS problem using the underlying sampler
        try:
            return self.sampler.sample(ahs_program)
        except Exception as e:
            print(f"An error occurred while sampling: {e}")
            raise

[Aaron-Robertson]

@golanor Not sure if you're still interested, but I reached out to Quera yesterday about a broken link in the optimization section of their site, and turns out it's the exact solution to the problem at hand 🤣! Classic, anyways I'm reading through it[1] now, to incorporate properly with the above stubs.

If you are still interested in collaborating, we may want to hop on a call and discuss the path forward and which pieces to each take on. Feel free to message me on the UF discord if that's easy, I'm A-Aron.

[JordanAWS]

I'd be happy to work on this with @Aaron-Robertson if he wants, or to work on it alone.

I want to make sure I understand the scope of the work here:

1. Create a new class BraketAquilaSampler in the plugin.
2. Modify the examples to work on Aquila.
3. Test on the device.

Is that correct @JordanAWS ? Any other things to take into consideration? (Other than asking for credits for testing the examples)

Yep!

[JordanAWS]

@Aaron-Robertson It looks like there's still some work remaining on your PR. Do you think you'll be able to finish it by the end of UnitaryHack today?