Merge pull request #62 from JuliaQUBO/feature/issue-61-topology-plots

Add native Pegasus and Zephyr plot helpers

Author: bernalde

CondaPkg.toml

@@ -2,6 +2,7 @@ channels = ["anaconda", "conda-forge"]
 
 [deps]
 python = ">=3.11,<=3.13"
+matplotlib = ">=3"
 
 [pip.deps]
 dwave-ocean-sdk = "==9.3.0"

README.md

@@ -123,6 +123,27 @@ current `QUBODrivers` interface cleanly.
 ## API Token
 To use D-Wave's QPU it is necessary to obtain an API Token from [Leap](https://cloud.dwavesys.com/leap/).
 
+## Topology and Embedding Plots
+`DWave.draw_topology` and `DWave.draw_embedding` wrap D-Wave NetworkX's
+Pegasus and Zephyr plotting helpers and return a Matplotlib figure by default.
+Both helpers accept `DWave.WorkingGraph` values built from sampler metadata, so
+plots use the full calibrated working graph rather than only the embedded
+qubits.
+
+```julia
+using DWave
+using QUBOTools
+
+metadata = QUBOTools.metadata(sampleset)
+
+DWave.draw_topology(DWave.WorkingGraph(metadata))
+DWave.draw_embedding(metadata)
+```
+
+Pass `ax = existing_axis` to draw into an existing Matplotlib axis; all other
+keyword arguments are forwarded to the corresponding D-Wave NetworkX draw
+function.
+
 ## Development Checks
 To verify that DWave can share a CondaPkg environment with the other
 Python-backed JuliaQUBO benchmark drivers, run:

src/wrapper/topology.jl

@@ -333,6 +333,169 @@ function _dnx_working_graph(arch::WorkingGraph)
     return nothing
 end
 
+function _dnx_hardware_graph(arch::Pegasus)
+    return _dnx().pegasus_graph(
+        arch.size;
+        node_list = _py(arch.nodes),
+        edge_list = _py(arch.edges),
+    )
+end
+
+function _dnx_hardware_graph(arch::Zephyr)
+    return _dnx().zephyr_graph(
+        arch.size,
+        arch.shore_size;
+        node_list = _py(arch.nodes),
+        edge_list = _py(arch.edges),
+    )
+end
+
+function _dnx_hardware_graph(arch::WorkingGraph)
+    graph = _dnx_working_graph(arch)
+
+    graph === nothing && throw(ArgumentError(
+        "D-Wave drawing helpers require Pegasus or Zephyr topology metadata with a nonempty shape",
+    ))
+
+    return graph
+end
+
+function _hardware_topology(source)
+    source isa DWaveHardwareTopology && return source
+    source isa AbstractDict && return WorkingGraph(source)
+    source isa QUBOTools.SampleSet && return WorkingGraph(QUBOTools.metadata(source))
+    source isa PythonCall.Py && return WorkingGraph(source)
+
+    throw(ArgumentError(
+        "expected a D-Wave hardware topology, metadata dictionary, sample set, or sampler",
+    ))
+end
+
+_draw_topology_function(::Pegasus) = _dnx().draw_pegasus
+_draw_topology_function(::Zephyr) = _dnx().draw_zephyr
+
+function _draw_topology_function(arch::WorkingGraph)
+    arch.topology_type == "pegasus" && return _dnx().draw_pegasus
+    arch.topology_type == "zephyr" && return _dnx().draw_zephyr
+
+    throw(ArgumentError("D-Wave drawing helpers only support Pegasus and Zephyr topologies"))
+end
+
+_draw_embedding_function(::Pegasus) = _dnx().draw_pegasus_embedding
+_draw_embedding_function(::Zephyr) = _dnx().draw_zephyr_embedding
+
+function _draw_embedding_function(arch::WorkingGraph)
+    arch.topology_type == "pegasus" && return _dnx().draw_pegasus_embedding
+    arch.topology_type == "zephyr" && return _dnx().draw_zephyr_embedding
+
+    throw(ArgumentError("D-Wave embedding drawing helpers only support Pegasus and Zephyr topologies"))
+end
+
+function _matplotlib_figure_axes()
+    figure_axes = PythonCall.pyimport("matplotlib.pyplot").subplots()
+
+    return figure_axes[0], figure_axes[1]
+end
+
+function _draw_result(draw_function, graph, args...; ax = nothing, kwargs...)
+    if ax === nothing
+        figure, draw_axis = _matplotlib_figure_axes()
+        draw_function(graph, args...; ax = draw_axis, kwargs...)
+
+        return figure
+    else
+        draw_function(graph, args...; ax = ax, kwargs...)
+
+        return ax
+    end
+end
+
+function _normalise_embedding_for_drawing(embedding_data)
+    normal_embedding = _normalise_embedding(embedding_data)
+
+    normal_embedding === nothing && throw(ArgumentError(
+        "embedding must be a dictionary mapping variables to integer qubit chains",
+    ))
+
+    return normal_embedding
+end
+
+function _embedding_from_metadata(metadata::AbstractDict)
+    normal_embedding = embedding(metadata)
+
+    normal_embedding === nothing && throw(ArgumentError(
+        "metadata does not contain a D-Wave embedding; sample with return_embedding=true",
+    ))
+
+    return normal_embedding
+end
+
+@doc raw"""
+    DWave.draw_topology(source; kwargs...)
+
+Draw a Pegasus or Zephyr hardware topology with D-Wave NetworkX.
+
+`source` may be a `Pegasus`, `Zephyr`, `WorkingGraph`, D-Wave sampler,
+`QUBOTools.SampleSet`, or metadata dictionary accepted by `WorkingGraph`.
+Keyword arguments are forwarded to `dwave_networkx.draw_pegasus` or
+`dwave_networkx.draw_zephyr`.
+
+By default this creates and returns a Matplotlib figure, which notebooks can
+display directly. If `ax` is supplied, drawing is performed on that axis and
+the same axis is returned.
+"""
+function draw_topology(arch::DWaveHardwareTopology; kwargs...)
+    return _draw_result(
+        _draw_topology_function(arch),
+        _dnx_hardware_graph(arch);
+        kwargs...,
+    )
+end
+
+function draw_topology(source; kwargs...)
+    return draw_topology(_hardware_topology(source); kwargs...)
+end
+
+@doc raw"""
+    DWave.draw_embedding(source, embedding; kwargs...)
+    DWave.draw_embedding(sampleset_or_metadata; kwargs...)
+
+Draw a returned minor embedding over the full Pegasus or Zephyr working graph
+with D-Wave NetworkX.
+
+The two-argument form accepts any topology source supported by
+`draw_topology` plus an embedding dictionary of the form
+`Dict{Int,Vector{Int}}`. The one-argument form extracts both the working graph
+and embedding from a `QUBOTools.SampleSet` or metadata dictionary, preserving
+compatibility with `DWave.WorkingGraph(QUBOTools.metadata(sampleset))`.
+
+Keyword arguments are forwarded to `dwave_networkx.draw_pegasus_embedding` or
+`dwave_networkx.draw_zephyr_embedding`. By default this creates and returns a
+Matplotlib figure; if `ax` is supplied, that axis is returned.
+"""
+function draw_embedding(arch::DWaveHardwareTopology, embedding_data::AbstractDict; kwargs...)
+    return _draw_result(
+        _draw_embedding_function(arch),
+        _dnx_hardware_graph(arch),
+        _py(_normalise_embedding_for_drawing(embedding_data));
+        kwargs...,
+    )
+end
+
+function draw_embedding(source, embedding_data::AbstractDict; kwargs...)
+    return draw_embedding(_hardware_topology(source), embedding_data; kwargs...)
+end
+
+function draw_embedding(metadata::AbstractDict; kwargs...)
+    return draw_embedding(WorkingGraph(metadata), _embedding_from_metadata(metadata); kwargs...)
+end
+
+function draw_embedding(sampleset::QUBOTools.SampleSet; kwargs...)
+    metadata = QUBOTools.metadata(sampleset)
+
+    return draw_embedding(WorkingGraph(metadata), _embedding_from_metadata(metadata); kwargs...)
+end
+
 function _dnx_layout_points(layout, nodes::Vector{Int})
     points = Vector{QUBOTools.Point{2,Float64}}(undef, length(nodes))
 

test/topology.jl

@@ -5,6 +5,7 @@ import QUBOTools
 import Test
 
 const Graphs = DWave.Graphs
+const PythonCall = DWave.PythonCall
 
 function _node_edge_pairs(arch)
     graph = QUBOTools.topology(QUBOTools.layout(arch))
@@ -23,6 +24,30 @@ function _point_count(points)
     return length(unique(points))
 end
 
+function _py_len(value)
+    return PythonCall.pyconvert(Int, PythonCall.pybuiltins.len(value))
+end
+
+function _py_class_name(value)
+    return PythonCall.pyconvert(String, value.__class__.__name__)
+end
+
+function _close_figure(figure)
+    PythonCall.pyimport("matplotlib.pyplot").close(figure)
+
+    return nothing
+end
+
+function _use_matplotlib_test_backend()
+    config_dir = get!(ENV, "MPLCONFIGDIR") do
+        mktempdir()
+    end
+    PythonCall.pyimport("os").environ["MPLCONFIGDIR"] = config_dir
+    PythonCall.pyimport("matplotlib").use("Agg"; force = true)
+
+    return nothing
+end
+
 Test.@testset "Pegasus layout uses Ocean topology data" begin
     arch = DWave.Pegasus(2)
     arch_layout = QUBOTools.layout(arch)
@@ -113,6 +138,86 @@ Test.@testset "WorkingGraph unknown topology uses graph geometry fallback" begin
     Test.@test _point_count(points) == 3
 end
 
+Test.@testset "D-Wave NetworkX drawing helpers use full working graphs" begin
+    _use_matplotlib_test_backend()
+
+    pegasus_metadata = Dict{String,Any}(
+        "dwave_info" => Dict{String,Any}(
+            "chip_info" => Dict{String,Any}(
+                "topology" => Dict{String,Any}("type" => "pegasus", "shape" => Any[2]),
+                "qubits" => Any[2, 3, 28],
+                "couplers" => Any[Any[2, 3], Any[2, 28]],
+            ),
+            "embedding_context" => Dict{String,Any}(
+                "embedding" => Dict{Any,Any}(1 => Any[2], 2 => Any[28]),
+            ),
+        ),
+    )
+    pegasus = DWave.WorkingGraph(pegasus_metadata)
+    pegasus_graph = DWave._dnx_hardware_graph(pegasus)
+
+    Test.@test isdefined(DWave, :draw_topology)
+    Test.@test isdefined(DWave, :draw_embedding)
+    Test.@test _py_len(pegasus_graph.nodes()) == length(pegasus.nodes)
+    Test.@test _py_len(pegasus_graph.edges()) == length(pegasus.edges)
+
+    figure = DWave.draw_topology(pegasus; node_size = 8, with_labels = false)
+
+    try
+        Test.@test _py_class_name(figure) == "Figure"
+    finally
+        _close_figure(figure)
+    end
+
+    figure = DWave.draw_embedding(pegasus_metadata; node_size = 8, with_labels = false)
+
+    try
+        Test.@test _py_class_name(figure) == "Figure"
+    finally
+        _close_figure(figure)
+    end
+
+    zephyr_metadata = Dict{String,Any}(
+        "dwave_info" => Dict{String,Any}(
+            "chip_info" => Dict{String,Any}(
+                "topology" => Dict{String,Any}("type" => "zephyr", "shape" => Any[2, 4]),
+                "qubits" => Any[0, 1, 4],
+                "couplers" => Any[Any[0, 1], Any[1, 4]],
+            ),
+            "embedding_context" => Dict{String,Any}(
+                "embedding" => Dict{Any,Any}(1 => Any[0], 2 => Any[4]),
+            ),
+        ),
+    )
+    zephyr = DWave.WorkingGraph(zephyr_metadata)
+    zephyr_graph = DWave._dnx_hardware_graph(zephyr)
+
+    Test.@test _py_len(zephyr_graph.nodes()) == length(zephyr.nodes)
+    Test.@test _py_len(zephyr_graph.edges()) == length(zephyr.edges)
+
+    figure = DWave.draw_topology(zephyr; node_size = 8, with_labels = false)
+
+    try
+        Test.@test _py_class_name(figure) == "Figure"
+    finally
+        _close_figure(figure)
+    end
+
+    figure = DWave.draw_embedding(zephyr, Dict(1 => [0], 2 => [4]); node_size = 8, with_labels = false)
+
+    try
+        Test.@test _py_class_name(figure) == "Figure"
+    finally
+        _close_figure(figure)
+    end
+
+    Test.@test_throws ArgumentError DWave.draw_embedding(Dict{String,Any}(
+        "dwave_info" => Dict{String,Any}(
+            "chip_info" => pegasus_metadata["dwave_info"]["chip_info"],
+        ),
+    ))
+end
+
 if DWave.__auth__(; verbose = false)
     Test.@testset "WorkingGraph builds from live D-Wave sampler metadata" begin
         sampler = DWave.dwave_system.DWaveSampler(; token = ENV["DWAVE_API_TOKEN"])