Merge pull request #59 from JuliaQUBO/fix/issue-58-return-embeddings

Expose returned embeddings for supplied samplers

Author: bernalde

src/sampler.jl

@@ -104,6 +104,81 @@ function _dwave_chip_info(dwave_sampler)
     return chip_info
 end
 
+function _normalise_embedding_index(value)
+    value isa Integer && return Int(value)
+
+    if value isa AbstractString
+        index = tryparse(Int, value)
+        index === nothing || return index
+    end
+
+    return nothing
+end
+
+function _normalise_embedding_chain(chain)
+    chain isa AbstractVector || chain isa Tuple || return nothing
+
+    normal_chain = Int[]
+    sizehint!(normal_chain, length(chain))
+
+    for qubit in chain
+        index = _normalise_embedding_index(qubit)
+        index === nothing && return nothing
+
+        push!(normal_chain, index)
+    end
+
+    return normal_chain
+end
+
+function _normalise_embedding(embedding)
+    embedding isa AbstractDict || return nothing
+
+    normal_embedding = Dict{Int,Vector{Int}}()
+
+    for (variable, chain) in pairs(embedding)
+        normal_variable = _normalise_embedding_index(variable)
+        normal_variable === nothing && return nothing
+        normal_chain = _normalise_embedding_chain(chain)
+        normal_chain === nothing && return nothing
+
+        normal_embedding[normal_variable] = normal_chain
+    end
+
+    return normal_embedding
+end
+
+function _normalise_dwave_embedding!(dwave_info::AbstractDict)
+    context = get(dwave_info, "embedding_context", nothing)
+    context isa AbstractDict || return dwave_info
+
+    normal_embedding = _normalise_embedding(get(context, "embedding", nothing))
+    normal_embedding === nothing || (context["embedding"] = normal_embedding)
+
+    return dwave_info
+end
+
+@doc raw"""
+    DWave.embedding(sampleset_or_metadata)
+
+Return the minor embedding recorded in D-Wave sample-set metadata, or `nothing`
+when no embedding was returned. Embeddings are normalized as
+`Dict{Int,Vector{Int}}`.
+"""
+function embedding(metadata::AbstractDict)
+    dwave_info = haskey(metadata, "dwave_info") ? metadata["dwave_info"] : metadata
+    dwave_info isa AbstractDict || return nothing
+
+    context = get(dwave_info, "embedding_context", nothing)
+    context isa AbstractDict || return nothing
+
+    return _normalise_embedding(get(context, "embedding", nothing))
+end
+
+function embedding(sampleset::QUBOTools.SampleSet)
+    return embedding(QUBOTools.metadata(sampleset))
+end
+
 function QUBODrivers.sample(sampler::Optimizer{T}) where {T}
     # Ising Model
     n, h, J, α, β = QUBOTools.ising(sampler, :dict; sense = :min)
@@ -112,8 +187,9 @@ function QUBODrivers.sample(sampler::Optimizer{T}) where {T}
     num_reads = MOI.get(sampler, DWave.NumberOfReads())
     final_num_reads = MOI.get(sampler, QUBODrivers.FinalNumberOfReads())
     sample_params = Dict{Symbol,Any}(
-        :num_reads      => final_num_reads,
-        :annealing_time => MOI.get(sampler, DWave.AnnealingTime()),
+        :num_reads         => final_num_reads,
+        :annealing_time    => MOI.get(sampler, DWave.AnnealingTime()),
+        :return_embedding => MOI.get(sampler, DWave.ReturnEmbedding()),
     )
     dwave_sampler = MOI.get(sampler, DWave.Sampler())
 
@@ -123,15 +199,14 @@ function QUBODrivers.sample(sampler::Optimizer{T}) where {T}
                 token = get(ENV, "DWAVE_API_TOKEN", nothing)
             )
         )
-
-        sample_params[:return_embedding] = MOI.get(sampler, DWave.ReturnEmbedding())
     end
 
     # Results
     samples = QUBOTools.Sample{T,Int}[]
     results = @timed dwave_sampler.sample_ising(h, J; sample_params...)
     var_map = pyconvert.(Int, [var for var in results.value.variables])
     dw_info = jl_object(results.value.info)
+    _normalise_dwave_embedding!(dw_info)
     chip_info = _dwave_chip_info(dwave_sampler)
 
     if !isempty(chip_info)

test/dwave_metadata.jl

@@ -45,7 +45,54 @@ sampler = FakeEmbeddingSampler()
     return ans.sampler
 end
 
-function _wrapper_dwave_sampleset(sampler)
+function _fake_return_embedding_sampler()
+    ans = DWave.PythonCall.pyexec(
+        @NamedTuple{sampler::DWave.PythonCall.Py,calls::DWave.PythonCall.Py},
+        """
+class FakeEmbeddingSampler:
+    def __init__(self):
+        self.calls = []
+        self.child = type("FakeChildSampler", (), {})()
+        self.child.properties = {
+            "chip_id": "mock-chip",
+            "topology": {"type": "pegasus", "shape": [16]},
+            "category": "qpu",
+            "qubits": [3, 4, 5],
+            "couplers": [[3, 4], [4, 5]],
+            "num_qubits": 3,
+        }
+        self.child.solver = type("FakeSolver", (), {"name": "mock-solver"})()
+
+    def sample_ising(self, h, J, **params):
+        self.calls.append(params)
+        info = {
+            "problem_id": "mock-problem",
+            "timing": {"qpu_access_time": 42},
+        }
+        if params.get("return_embedding"):
+            info["embedding_context"] = {"embedding": {1: (3, 4), 2: (5,)}}
+
+        return type(
+            "FakeSampleSet",
+            (),
+            {
+                "variables": [1, 2],
+                "record": [([1, -1], -1.25, 3)],
+                "info": info,
+            },
+        )()
+
+sampler = FakeEmbeddingSampler()
+calls = sampler.calls
+""",
+        @__MODULE__,
+        (),
+    )
+
+    return ans
+end
+
+function _wrapper_dwave_sampleset(sampler; return_embedding::Bool = false)
     model = MOI.instantiate(DWave.Optimizer; with_bridge_type = Float64)
     variables, _ = MOI.add_constrained_variables(model, fill(QUBODrivers.Spin(), 2))
 
@@ -65,6 +112,7 @@ function _wrapper_dwave_sampleset(sampler)
         ),
     )
     MOI.set(model, MOI.RawOptimizerAttribute("sampler"), sampler)
+    MOI.set(model, MOI.RawOptimizerAttribute("return_embedding"), return_embedding)
 
     MOI.optimize!(model)
 
@@ -100,3 +148,16 @@ Test.@testset "DWave metadata includes chip info" begin
     Test.@test chip_info["couplers"] == Any[Any[0, 1], Any[1, 4]]
     Test.@test chip_info["num_qubits"] == 3
 end
+
+Test.@testset "DWave metadata exposes returned embeddings for supplied samplers" begin
+    ans = _fake_return_embedding_sampler()
+    sampleset = _wrapper_dwave_sampleset(ans.sampler; return_embedding = true)
+    metadata = QUBOTools.metadata(sampleset)
+    calls = DWave.jl_object(ans.calls)
+
+    Test.@test calls[1]["return_embedding"] == true
+    Test.@test DWave.embedding(sampleset) == Dict(1 => [3, 4], 2 => [5])
+    Test.@test DWave.embedding(metadata) == Dict(1 => [3, 4], 2 => [5])
+    Test.@test metadata["dwave_info"]["embedding_context"]["embedding"] ==
+        Dict(1 => [3, 4], 2 => [5])
+end