Ocean_Composites.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Exploring Ocean Composites\n",
    "\n",
    "[Composites](https://docs.ocean.dwavesys.com/en/stable/concepts/samplers.html#composites) play a major role in the Ocean SDK and allow users to solve problems in [multiple steps](https://en.wikipedia.org/wiki/Composite_pattern). \n",
    "Ocean composites may provide pre-processing, post-processing, or other functionalities that could be useful when composing a [sampler](https://docs.ocean.dwavesys.com/en/stable/concepts/samplers.html#samplers).\n",
    "\n",
    "Note: As it stands, composites can only be used to compose samplers for [Binary Quadratic Models (BQMs)]( https://docs.ocean.dwavesys.com/en/stable/concepts/bqm.html#bqm-sdk).\n",
    "\n",
    "In this notebook, we cover several composites, exploring how they can be used and why they are useful:\n",
    "\n",
    "1. [StructureComposite](https://docs.ocean.dwavesys.com/en/stable/docs_dimod/reference/sampler_composites/composites.html#module-dimod.reference.composites.structure)\n",
    "2. [EmbeddingComposite](https://docs.ocean.dwavesys.com/en/stable/docs_system/reference/composites.html#embeddingcomposite)\n",
    "3. [FixedEmbeddingComposite](https://docs.ocean.dwavesys.com/en/stable/docs_system/reference/composites.html#fixedembeddingcomposite)\n",
    "4. [TruncateComposite](https://docs.ocean.dwavesys.com/en/stable/docs_dimod/reference/sampler_composites/composites.html#module-dimod.reference.composites.truncatecomposite)\n",
    "5. [FixVariablesComposite](https://docs.ocean.dwavesys.com/en/stable/docs_preprocessing/reference/composites.html#fix-variables-composite)\n",
    "6. [TrackingComposite](https://docs.ocean.dwavesys.com/en/stable/docs_dimod/reference/sampler_composites/composites.html#module-dimod.reference.composites.tracking)\n",
    "6. [SteepestDescentComposite](https://docs.ocean.dwavesys.com/en/stable/docs_greedy/reference/composites.html#steepestdescentcomposite)\n",
    "\n",
    "Note: One sampler in this notebook, the [DWaveSampler](https://docs.ocean.dwavesys.com/projects/system/en/stable/reference/samplers.html#dwavesampler), uses a Quantum Processing Unit (QPU). Usage of QPU requires a Leap account with access to QPU solver(s)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1) Structure Composite\n",
    "\n",
    "The `StructureComposite` can be used to compose a sampler that targets a particular graph structure. The ability to do this is useful for simulating proposed working graphs of QPU designs. In this demonstration, we will show how to build a composed sampler using the Simulated Annealing Sampler and Structure Composite, targeting a 3-dimensional cubic lattice structure. \n",
    "\n",
    "Lattice structures are particularly interesting due to the variety of applications in the sciences. To name a few:\n",
    "- Kosterlitz–Thouless phase transition analysis (https://www.science.org/doi/10.1126/science.aat2025, https://www.nature.com/articles/s41586-018-0410-x)\n",
    "- Lattice-based cryptography (https://en.wikipedia.org/wiki/Lattice-based_cryptography)\n",
    "- Study of crystalline structures (https://en.wikipedia.org/wiki/Crystal_structure)\n",
    "\n",
    "First, let's use NetworkX to build a lattice-structured graph for our composed simple 3d lattice sampler, and view the lattice in a 3d plot with matplotlib."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dimod\n",
    "import networkx as nx\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits import mplot3d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build a simple 3D lattice with grid_graph\n",
    "\n",
    "lattice = nx.grid_graph(dim=(4,4,4))\n",
    "\n",
    "%matplotlib notebook\n",
    "\n",
    "nodes_xyz = np.array(lattice.nodes)\n",
    "edges_xyz = np.array(lattice.edges)\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = plt.axes(projection=\"3d\")\n",
    "ax.set(xlabel=\"X\", ylabel=\"Y\", zlabel=\"Z\",title=\"4x4x4 Simple Lattice\")\n",
    "\n",
    "#plot points\n",
    "ax.scatter3D(nodes_xyz.T[0], nodes_xyz.T[1], nodes_xyz.T[2], c='b')\n",
    "\n",
    "#plot the edges\n",
    "for edge in edges_xyz:\n",
    "    ax.plot3D(*zip(edge[0], edge[1]), color=\"orchid\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build a composed structured sampler\n",
    "\n",
    "Now we are ready to build our composed structured sampler with our lattice graph.\n",
    "\n",
    "Let's look at the signature of the `StructureComposite` class, found in `dimod`, to see what we need:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "?dimod.StructureComposite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We provide a base sampler, in this case the `SimulatedAnnealingSampler` from `dwave-samplers`, a list of nodes, and a list of edges.\n",
    "The edge and node lists will come from our lattice graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dwave.samplers import SimulatedAnnealingSampler\n",
    "\n",
    "sampler = SimulatedAnnealingSampler()\n",
    "\n",
    "structured_sampler = dimod.StructureComposite(sampler, list(lattice.nodes), list(lattice.edges))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With the structured sampler ready, we can sample from BQMs with the same target graph.\n",
    "Next, let's generate a random BQM that has the exact same graph as our structured sampler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build a BQM with random biases and same structure as our structured sampler.\n",
    "\n",
    "linear = {node:np.random.randint(-20,20) for node in lattice.nodes}\n",
    "quadratic = {edge:np.random.randint(-20,20) for edge in lattice.edges}\n",
    "\n",
    "bqm = dimod.BQM(linear, quadratic, \"BINARY\")\n",
    "\n",
    "response_structured = structured_sampler.sample(bqm, num_reads=10)\n",
    "\n",
    "print(response_structured)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: We could sample from a lattice graph without the structured sampler and just use the plain `SimulatedAnnealingSampler`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2) Embedding Composite\n",
    "\n",
    "This is great! However, we are restricted to BQMs with a very specfic graph structure - one exactly matching the structured sampler. We can, however, exploit this structure in a meaningful way by composing our sampler with yet another composite - the `EmbeddingComposite`, found in the package `dwave-system`.\n",
    "\n",
    "Looking at the signature tells us that we need to provide a *structured* sampler, among other optional arguments.\n",
    "We require the input sampler to be structured, and it is worth mentioning that every QPU has a structure (Chimera, Pegasus are names of working graphs).\n",
    "\n",
    "Our composed sampler has a lattice structure, so we can compose it with the `EmbeddingComposite` and embed problem graphs that are embeddable. This has the benefit of solving problems that don't have the exact same structure as the sampler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dwave.system import EmbeddingComposite\n",
    "\n",
    "?EmbeddingComposite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, there may be many graphs that we can map to a lattice. Today, we thought it would be interesting and perhaps motivating to show embedding the same structure of lattice, just of a smaller size. This motivates thinking of parallelization via structure.\n",
    "\n",
    "First, let's build our composed embedding sampler, that looks to embed the input problem's graph into our original lattice graph.\n",
    "Then, let's create a smaller 3d lattice to sample with the composed sampler."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embedding_sampler = EmbeddingComposite(structured_sampler)\n",
    "\n",
    "# Create a smaller simple 3d lattice\n",
    "\n",
    "lattice_small = nx.grid_graph(dim=(2,2,2))\n",
    "\n",
    "linear_small = {node:np.random.randint(-20,20) for node in lattice_small.nodes}\n",
    "quadratic_small = {edge:np.random.randint(-20,20) for edge in lattice_small.edges}\n",
    "\n",
    "small_bqm = dimod.BQM(linear_small, quadratic_small, \"BINARY\")\n",
    "\n",
    "response_embedded = embedding_sampler.sample(small_bqm, num_reads=10, return_embedding=True)\n",
    "\n",
    "print(response_embedded)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's visualize the embedding of the smaller lattice into the larger one:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_embedding = []\n",
    "for embedding in response_embedded.info['embedding_context']['embedding'].values():\n",
    "    node_embedding += [*embedding]\n",
    "    \n",
    "node_embedding = np.array(node_embedding)\n",
    "\n",
    "%matplotlib notebook\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = plt.axes(projection=\"3d\")\n",
    "ax.set(xlabel=\"X\", ylabel=\"Y\", zlabel=\"Z\", title=\"2x2x2 Lattice Embedding\")\n",
    "\n",
    "#plot points\n",
    "ax.scatter3D(nodes_xyz.T[0], nodes_xyz.T[1], nodes_xyz.T[2], c='b')\n",
    "ax.scatter3D(node_embedding.T[0], node_embedding.T[1], node_embedding.T[2], s=50, c='r')\n",
    "\n",
    "#plot the edges\n",
    "for edge in edges_xyz:\n",
    "    if any(np.array_equal(edge[0], x) for x in node_embedding) and any(np.array_equal(edge[1], x) for x in node_embedding):\n",
    "        ax.plot3D(*zip(edge[0], edge[1]), \"r\")\n",
    "    else:\n",
    "        ax.plot3D(*zip(edge[0], edge[1]), color=\"orchid\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Embedding with Chains\n",
    "\n",
    "Sometimes, we need to embed a problem graph that could have a different topology, or a similar structure but with higher connectivity on average. In these situations, we could end up mapping a node in our problem graph to several nodes in our structured sampler's graph to account for the different topologies. This creates \"chains\" of nodes and edges that are used to represent a node in our problem graph via embedding.\n",
    "\n",
    "Let's try to embed a 3d lattice that has periodic boundary conditions that result in a graph with a different topology than our simple lattice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a 3x2x2 lattice with periodic boundary conditions\n",
    "\n",
    "periodic_lattice = nx.grid_graph(dim=(3,2,2), periodic=True)\n",
    "\n",
    "linear_periodic = {node:np.random.randint(-20,20) for node in periodic_lattice.nodes}\n",
    "quadratic_periodic = {edge:np.random.randint(-20,20) for edge in periodic_lattice.edges}\n",
    "\n",
    "periodic_lattice_bqm = dimod.BQM(linear_periodic, quadratic_periodic, \"BINARY\")\n",
    "\n",
    "response_embedded = embedding_sampler.sample(periodic_lattice_bqm, num_reads=10, return_embedding=True)\n",
    "\n",
    "print(response_embedded)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response_embedded.info['embedding_context']['embedding']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Above, we see that some of the nodes from our input graph are mapped to sometimes several nodes in our structured sampler's graph. We can take a look at the node degrees to get a sense of the topological difference in these graphs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Periodic lattice node degrees:\")\n",
    "print(periodic_lattice.degree(), \"\\n\")\n",
    "print(\"4x4x4 simple lattice node degrees:\")\n",
    "print(lattice.degree())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The visualization becomes pretty messy, but let's take a look at how the embedding worked out in this more complex case.\n",
    "We end up mapping single nodes in the `periodic_lattice_bqm` graph to multiple nodes in our structured sampler's graph due to the periodic boundary conditions. These are known as chains in embedding; using multiple nodes and edges to represent one node."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "node_embedding = []\n",
    "for embedding in response_embedded.info['embedding_context']['embedding'].values():\n",
    "    node_embedding += [*embedding]\n",
    "    \n",
    "node_embedding = np.array(node_embedding)\n",
    "\n",
    "%matplotlib notebook\n",
    "\n",
    "fig = plt.figure()\n",
    "ax = plt.axes(projection=\"3d\")\n",
    "ax.set(xlabel=\"X\", ylabel=\"Y\", zlabel=\"Z\", title=\"3x2x2 Periodic Lattice Embedding - Chains\")\n",
    "\n",
    "#plot points\n",
    "ax.scatter3D(nodes_xyz.T[0], nodes_xyz.T[1], nodes_xyz.T[2], c='b')\n",
    "ax.scatter3D(node_embedding.T[0], node_embedding.T[1], node_embedding.T[2], s=50, c='r')\n",
    "\n",
    "#plot the edges\n",
    "for edge in edges_xyz:\n",
    "    if any(np.array_equal(edge[0], x) for x in node_embedding) and any(np.array_equal(edge[1], x) for x in node_embedding):\n",
    "        ax.plot3D(*zip(edge[0], edge[1]), \"r\")\n",
    "    else:\n",
    "        ax.plot3D(*zip(edge[0], edge[1]), color=\"orchid\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3) FixedEmbeddingComposite\n",
    "\n",
    "This is especially important if your workflow involves solving a problem with a repeating structure, and you are worried about performace. Embedding isn't a cheap operation on complex problems, it definitely takes some time. Let's take a look at how we can boost the performance of the sampler if we determine an embedding to our test problem in advance, using the `FixedEmbeddingComposite` found in `dwave-system`.\n",
    "\n",
    "Looking at the signature we see that we need a structured sampler, as well as to provide an embedding dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dwave.system import FixedEmbeddingComposite\n",
    "?FixedEmbeddingComposite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We start by constructing a larger lattice based sampler\n",
    "\n",
    "large_lattice = nx.grid_graph(dim=(10,10,10))\n",
    "test_structured_sampler = dimod.StructureComposite(SimulatedAnnealingSampler(), list(large_lattice.nodes), list(large_lattice.edges))\n",
    "\n",
    "# A test lattice to build a BQM from\n",
    "test_lattice = nx.grid_graph(dim=(4,4,4), periodic=True)\n",
    "test_linear = {node:np.random.randint(-50,50) for node in test_lattice.nodes}\n",
    "test_quadratic = {edge:np.random.randint(-50,50) for edge in test_lattice.edges}\n",
    "\n",
    "test_bqm = dimod.BQM(test_linear, test_quadratic, 'BINARY')\n",
    "\n",
    "# And for comparison, composed samplers from both Embedding and FixedEmbedding composites\n",
    "# We need to retrieve the embedding from a sampler run to build our FixedEmbedding sampler.\n",
    "\n",
    "test_embedding_sampler = EmbeddingComposite(test_structured_sampler)\n",
    "\n",
    "test_response = test_embedding_sampler.sample(test_bqm, num_reads=10, return_embedding=True)\n",
    "\n",
    "fixed_embedding_sampler = FixedEmbeddingComposite(test_structured_sampler,\n",
    "                                                  embedding=test_response.info['embedding_context']['embedding'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We are now ready to compare the performance of each sampler in producing a response.\n",
    "One has to find an embedding everytime (embedding_sampler), the other runs on the fixed embedding provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit test_embedding_sampler.sample(test_bqm, num_reads=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit fixed_embedding_sampler.sample(test_bqm, num_reads=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4) TruncateComposite\n",
    "\n",
    "To cap it off, let's take a look at another composite that acts as a post-processing layer. The `TruncateComposite` found in `dimod` allows us to filter the returned sample set, for example if we only care about a few of the lowest energy samples returned.\n",
    "\n",
    "The signature shows us we need to provide a sampler and number of rows to return, with some other optional arguments."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "?dimod.TruncateComposite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "truncated_sampler = dimod.TruncateComposite(embedding_sampler, 5)\n",
    "\n",
    "response_truncated = truncated_sampler.sample(small_bqm, num_reads=10)\n",
    "\n",
    "print(response_truncated)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's look at ways we can improve on our results."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5) FixVariablesComposite\n",
    "\n",
    "The `FixVariablesComposite` allows users to 'fix' some variables of a BQM before it is passed to a sampler, guaranteeing that returned samples will include the fixed variables. \n",
    "\n",
    "By fixing variables, we are reducing the size of the problem that actually needs to be solved. This is especially helpful for larger problems that cannot be easily solved on the QPU. \n",
    "\n",
    "The `FixVariablesComposite` can be found in `dwave-preprocessing`, where many of Ocean's pre-processing composites are located."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dwave.preprocessing import FixVariablesComposite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is its class signature:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "?FixVariablesComposite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `FixVariablesComposite` requires the user to specify a child sampler and the algorithm that determines which variables to fix. \n",
    "\n",
    "We can demonstrate this with a toy BQM of one ising variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dimod\n",
    "bqm = dimod.BinaryQuadraticModel.from_ising({'a':1}, {})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We'll use the `ExactSolver` from `dimod` as the child sampler, which is great for testing small problems because it calculates the energy for every possible sample."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "exact_solver = dimod.ExactSolver()\n",
    "sampler = FixVariablesComposite(exact_solver, algorithm='explicit')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, `algorithm` is set to `explicit`, meaning that the user must explicitly pass in a `fixed_variables` dict when sampling, as seen below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using only the ExactSolver\n",
    "sampleset = exact_solver.sample(bqm)\n",
    "print(\"Using only the ExactSolver\")\n",
    "print(sampleset)\n",
    "\n",
    "# Using the FixVariablesComposite and fixing the one variable to -1\n",
    "sampleset = sampler.sample(bqm, fixed_variables={'a':-1})\n",
    "print(\"\\nAlso using the FixVariablesComposite\")\n",
    "print(sampleset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of passing in the fixed variables explicitly, `algorithm` can also be set to `roof_duality`. With this algorithm, the composite will try to find minimizing assignments for some or all of the BQM's variables. For more information on the roof duality algorithm, see [`roof_duality()`](https://docs.ocean.dwavesys.com/en/stable/docs_preprocessing/reference/lower_bounds.html#dwave.preprocessing.lower_bounds.roof_duality)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampler = FixVariablesComposite(exact_solver, algorithm='roof_duality')\n",
    "sampleset = sampler.sample(bqm)\n",
    "print(sampleset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, we run the roof duality algorithm in `strict` mode. This means that we only fix variables when the variable assignments are True for all ground states. To demonstrate this, let's pick a different BQM:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bqm = dimod.BinaryQuadraticModel.from_ising({}, {'ab':-1})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This BQM has two ground states, either both `a` and `b` are -1 or both `a` and `b` are +1, which means that there are no variables that have a single value for all ground states. Thus, when `strict=True`, we don't have any fixed variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampleset = sampler.sample(bqm, strict=True)\n",
    "print(sampleset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When `strict=False`, we also fix the variables with assignments that are True for some but not all ground states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sampleset = sampler.sample(bqm, strict=False)\n",
    "print(sampleset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's test the composite out on a larger problem. We go back to our lattices:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import dimod\n",
    "from dwave.samplers import SimulatedAnnealingSampler\n",
    "from dwave.preprocessing import FixVariablesComposite\n",
    "\n",
    "# Build a BQM from a 3D lattice\n",
    "lattice = nx.grid_graph(dim=(20,20,20))\n",
    "\n",
    "linear = {node:np.random.randint(-50,50) for node in lattice.nodes}\n",
    "quadratic = {edge:np.random.randint(-50,50) for edge in lattice.edges}\n",
    "\n",
    "bqm = dimod.BinaryQuadraticModel(linear, quadratic, \"BINARY\")\n",
    "\n",
    "num_reads = 10\n",
    "\n",
    "# Solve with SimulatedAnnealingSampler\n",
    "sampler = SimulatedAnnealingSampler()\n",
    "sampleset = sampler.sample(bqm, num_reads=num_reads)\n",
    "\n",
    "# Add the FixVariablesComposite\n",
    "fixed_sampler = FixVariablesComposite(sampler, algorithm='roof_duality')\n",
    "fixed_sampleset = fixed_sampler.sample(bqm, strict=True, num_reads=num_reads)\n",
    "\n",
    "# Compare results\n",
    "plt.figure().clear('all')\n",
    "plt.plot(list(range(num_reads)), sorted(sampleset.record.energy))\n",
    "plt.plot(list(range(num_reads)), sorted(fixed_sampleset.record.energy))\n",
    "plt.ylabel('Energy')\n",
    "plt.xlabel('Samples')\n",
    "plt.legend(['SimulatedAnnealingSampler', 'With FixVariablesComposite'], loc='upper left')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Although we found improvements in our example above, it is not guaranteed that the FixVariablesComposite will always find variables to fix, nor is it a guarantee that results will improve due to these fixes. In 'roof_duality' mode, it can be difficult for users to determine if the composite was able to fix any variables. In this situation, users should take a look at the `TrackingComposite`, as described below."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6) TrackingComposite\n",
    "\n",
    "The `TrackingComposite`, found in `dimod`, is a composite that tracks inputs and outputs for the purpose of debugging and testing. It is especially useful when used with samplers made up of multiple composites."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dimod import TrackingComposite"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "?TrackingComposite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As seen in its signature, the `TrackingComposite` takes in a child sampler and a boolean for the `copy` argument, which determines whether or not the tracked inputs and ouputs are copied before they are stored. The `copy` argument defaults to `False`, but may be useful if the child sampler mutates values.\n",
    "\n",
    "Let's take a look at how the composite works. We'll start with a small BQM for demonstration. Once again, we can make use of `dimod`'s `ExactSolver`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "small_bqm = dimod.BinaryQuadraticModel({'a': 1, 'b': -1}, {}, 0, 'BINARY')\n",
    "track_sampler = TrackingComposite(dimod.ExactSolver())\n",
    "sampleset = track_sampler.sample(small_bqm)\n",
    "print(sampleset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`TrackingComposite.input` will return the most recent input to any sampling method. In this example, we know the most recent input is `small_bqm`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(track_sampler.input)\n",
    "print(\"Input is the same as original BQM: \", track_sampler.input == dict(bqm=small_bqm))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`TrackingComposite.output` will return the most recent output to any sampling method. In this example, we know the most recent output should be equivalent to `sampleset`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(track_sampler.output)\n",
    "print(\"Output is the sampleset printed above: \", track_sampler.output == sampleset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For all inputs and outputs tracked by the composite, we can use `TrackingComposite.inputs` and `TrackingComposite.outputs` respectively.\n",
    "\n",
    "In our small example, the `TrackingComposite` is not very useful because we already know the input and output to begin with. Let's see how it can be used in a more complex problem. Previously, we mentioned that users may find it difficult to determine whether or not the `FixVariablesComposite` was able to fix any variables when using the `roof_duality` mode. We can see how the `TrackingComposite` may provide a solution.\n",
    "\n",
    "Let's re-create our sampler and problem BQM from the `FixVariablesComposite` section above, adding in our new composite as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "import numpy as np\n",
    "import dimod\n",
    "from dwave.samplers import SimulatedAnnealingSampler\n",
    "from dwave.preprocessing import FixVariablesComposite\n",
    "\n",
    "# Build a BQM from a 3D lattice\n",
    "lattice = nx.grid_graph(dim=(20,20,20))\n",
    "\n",
    "linear = {node:np.random.randint(-50,50) for node in lattice.nodes}\n",
    "quadratic = {edge:np.random.randint(-50,50) for edge in lattice.edges}\n",
    "\n",
    "bqm = dimod.BinaryQuadraticModel(linear, quadratic, \"BINARY\")\n",
    "\n",
    "num_reads = 10\n",
    "\n",
    "# Add the FixVariablesComposite to fix variables and the TrackingComposite to track input/output\n",
    "track_sampler = TrackingComposite(SimulatedAnnealingSampler())\n",
    "fixed_sampler = FixVariablesComposite(track_sampler, algorithm='roof_duality')\n",
    "fixed_sampleset = fixed_sampler.sample(bqm, strict=True, num_reads=num_reads)\n",
    "print(fixed_sampleset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If the input of the TrackingComposite is not the same as the original `bqm`, then we know that the FixVariablesComposite was able to fix some variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_bqm = track_sampler.input['bqm']\n",
    "print(\"Size of original BQM: \", len(bqm))\n",
    "print(\"Number of variables fixed: \", len(bqm) - len(input_bqm))\n",
    "print(\"Size of BQM that was actually solved: \", len(input_bqm))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7) SteepestDescentComposite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's move on to post-processing optimization, which can often help make local improvements. The `SteepestDescentComposite` runs a greedy local optimization on a problem by taking samples from a child sampler as initial states.\n",
    "\n",
    "The `SteepestDescentComposite` can be found in `dwave-samplers`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dwave.samplers import SteepestDescentComposite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is its class signature:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "?SteepestDescentComposite"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this example, let's use the `DWaveSampler` as our child sampler. Our problem graph is incompatible with the QPU architecture, so we'll also have to use the `EmbeddingComposite`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dwave.system import DWaveSampler, EmbeddingComposite\n",
    "\n",
    "sampler = EmbeddingComposite(DWaveSampler())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Embedding a lattice onto the QPU can take a while, so let's use a smaller 3D lattice this time around and increase the range of our biases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import networkx as nx\n",
    "import numpy as np\n",
    "import dimod\n",
    "\n",
    "# Build a BQM from a 3D lattice\n",
    "lattice = nx.grid_graph(dim=(5,5,5))\n",
    "\n",
    "linear = {node:np.random.randint(-100,100) for node in lattice.nodes}\n",
    "quadratic = {edge:np.random.randint(-100,100) for edge in lattice.edges}\n",
    "\n",
    "bqm = dimod.BinaryQuadraticModel(linear, quadratic, \"BINARY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's compare the results with post-processing and without."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dwave.samplers import SteepestDescentComposite\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "num_reads = 10\n",
    "\n",
    "# Solve with DWaveSampler\n",
    "sampleset = sampler.sample(bqm, num_reads=num_reads)\n",
    "print(sampleset)\n",
    "\n",
    "# Solve with the SteepestDescentComposite\n",
    "greedy_sampler = SteepestDescentComposite(sampler)\n",
    "greedy_sampleset = greedy_sampler.sample(bqm, num_reads=num_reads)\n",
    "print(greedy_sampleset)\n",
    "\n",
    "# Compare results\n",
    "plt.figure().clear('all')\n",
    "plt.plot(list(range(num_reads)), sorted(sampleset.record.energy))\n",
    "plt.plot(list(range(num_reads)),sorted(greedy_sampleset.record.energy))\n",
    "\n",
    "plt.ylabel('Energy')\n",
    "plt.xlabel('Samples')\n",
    "plt.legend(['DWaveSampler', 'With SteepestDescentComposite'], loc='upper left')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Honorable mentions\n",
    "[TilingComposite](https://docs.ocean.dwavesys.com/projects/system/en/stable/reference/composites.html#tilingcomposite): Tiles input problem graph across the QPU graph (Chimera, Pegasus). Enables parallel sampling for small problems.\n",
    "\n",
    "[ConnectedComponentsComposite](https://docs.ocean.dwavesys.com/en/latest/docs_preprocessing/reference/composites.html#connected-components-composite): Decomposes an input problem into its connected components, merges best samples of each reponse and returns a merged sample set."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}