Skip to content

[codex] Add QUBODrivers 0.6.1 conformance #27

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bernalde merged 1 commit into from
Jun 13, 2026
Merged

[codex] Add QUBODrivers 0.6.1 conformance #27

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,6 @@ QUBOTools = "60eb5b62-0a39-4ddc-84c5-97d2adff9319"
[compat]
MathOptInterface = "1"
PythonCall = "0.9.23"
QUBODrivers = "0.3, 0.4, 0.5, 0.6"
QUBOTools = "0.10, 0.11, 0.12, 0.13"
QUBODrivers = "0.6.1"
QUBOTools = "0.13"
julia = "1.10"
34 changes: 20 additions & 14 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ model = Model(PySA.Optimizer)
set_silent(model)
set_attribute(model, PySA.NumberOfReads(), 20)
set_attribute(model, PySA.NumberOfSweeps(), 64)
set_attribute(model, PySA.RandomSeed(), 1234)

n = 3
Q = [ -1 2 2
Expand All @@ -42,23 +43,28 @@ end

## Solver options

PySA.jl exposes the main PySA simulated annealing options as JuMP optimizer attributes:

| Attribute | PySA option | Default |
| --- | --- | --- |
| `PySA.NumberOfSweeps()` | `n_sweeps` | `32` |
| `PySA.NumberOfReplicas()` | `n_replicas` | `3` |
| `PySA.NumberOfReads()` | `n_reads` | `10` |
| `PySA.MinimumTemperature()` | `min_temp` | `1.0` |
| `PySA.MaximumTemperature()` | `max_temp` | `3.5` |
| `PySA.UpdateStrategy()` | `update_strategy` | `"sequential"` |
| `PySA.InitializeStrategy()` | `initialize_strategy` | `"ones"` |
| `PySA.RecomputeEnergy()` | `recompute_energy` | `false` |
| `PySA.SortOutputTemps()` | `sort_output_temps` | `true` |
| `PySA.Parallel()` | `parallel` | `true` |
PySA.jl exposes the main PySA simulated annealing options as JuMP optimizer attributes.
Each option can also be set with `MOI.RawOptimizerAttribute` using the raw key.
Legacy `n_*` raw keys remain supported for compatibility.

| Attribute | Raw key | Alias raw key | Default |
| --- | --- | --- | --- |
| `PySA.NumberOfSweeps()` | `num_sweeps` | `n_sweeps` | `32` |
| `PySA.NumberOfReplicas()` | `num_replicas` | `n_replicas` | `3` |
| `PySA.NumberOfReads()` | `num_reads` | `n_reads` | `10` |
| `PySA.FinalNumberOfReads()` | `final_num_reads` | - | `num_reads` |
| `PySA.RandomSeed()` | `seed` | - | `nothing` |
| `PySA.MinimumTemperature()` | `min_temp` | `minimum_temperature` | `1.0` |
| `PySA.MaximumTemperature()` | `max_temp` | `maximum_temperature` | `3.5` |
| `PySA.UpdateStrategy()` | `update_strategy` | - | `"sequential"` |
| `PySA.InitializeStrategy()` | `initialize_strategy` | - | `"ones"` |
| `PySA.RecomputeEnergy()` | `recompute_energy` | - | `false` |
| `PySA.SortOutputTemps()` | `sort_output_temps` | - | `true` |
| `PySA.Parallel()` | `parallel` | - | `true` |

Use `set_attribute(model, attribute, value)` to override these values before calling `optimize!`.
Use `set_silent(model)` to disable PySA solver output.
When `PySA.RandomSeed()` is set, PySA.jl runs the backend with `parallel=false` so NumPy and Numba random streams are reproducible.

**Note**: _The PySA wrapper for Julia is not officially supported by the National Aeronautics and Space Administration. If you are interested in official support for Julia from NASA, let them know!_

Expand Down
207 changes: 173 additions & 34 deletions src/PySA.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,20 +10,36 @@ import MathOptInterface as MOI
const np = PythonCall.pynew()
const pysa = PythonCall.pynew()
const pysa_sa = PythonCall.pynew()
const seed_numba_rng = PythonCall.pynew()

function __init__()
PythonCall.pycopy!(np, pyimport("numpy"))
PythonCall.pycopy!(pysa, pyimport("pysa"))
PythonCall.pycopy!(pysa_sa, pyimport("pysa.sa"))

ns = PythonCall.pydict()
pyimport("builtins").exec(
"""
import numpy as np
import numba

@numba.njit
def seed_numba_rng(seed):
np.random.seed(seed)
""",
ns,
)
PythonCall.pycopy!(seed_numba_rng, ns["seed_numba_rng"])
end

QUBODrivers.@setup Optimizer begin
name = "PySA"
version = v"0.1.0" # pysa version
attributes = begin
NumberOfSweeps["n_sweeps"]::Integer = 32
NumberOfReplicas["n_replicas"]::Integer = 3
NumberOfReads["n_reads"]::Integer = 10
NumberOfSweeps["num_sweeps"]::Integer = 32
NumberOfReplicas["num_replicas"]::Integer = 3
NumberOfReads["num_reads"]::Integer = 10
"seed"::Union{Integer,Nothing} = nothing
MinimumTemperature["min_temp"]::Float64 = 1.0
MaximumTemperature["max_temp"]::Float64 = 3.5
UpdateStrategy["update_strategy"]::String = "sequential"
Expand All @@ -34,6 +50,55 @@ QUBODrivers.@setup Optimizer begin
end
end

const RandomSeed = QUBODrivers.RandomSeed
const FinalNumberOfReads = QUBODrivers.FinalNumberOfReads

const _RawLegacyNumberOfSweeps = QUBODrivers.RawSamplerAttribute{Symbol("n_sweeps")}
const _RawLegacyNumberOfReplicas = QUBODrivers.RawSamplerAttribute{Symbol("n_replicas")}
const _RawLegacyNumberOfReads = QUBODrivers.RawSamplerAttribute{Symbol("n_reads")}
const _RawMinimumTemperature = QUBODrivers.RawSamplerAttribute{Symbol("minimum_temperature")}
const _RawMaximumTemperature = QUBODrivers.RawSamplerAttribute{Symbol("maximum_temperature")}

MOI.get(sampler::Optimizer, ::_RawLegacyNumberOfSweeps) = MOI.get(sampler, NumberOfSweeps())
MOI.get(sampler::Optimizer, ::_RawLegacyNumberOfReplicas) = MOI.get(sampler, NumberOfReplicas())
MOI.get(sampler::Optimizer, ::_RawLegacyNumberOfReads) = MOI.get(sampler, NumberOfReads())
MOI.get(sampler::Optimizer, ::_RawMinimumTemperature) = MOI.get(sampler, MinimumTemperature())
MOI.get(sampler::Optimizer, ::_RawMaximumTemperature) = MOI.get(sampler, MaximumTemperature())

function MOI.set(sampler::Optimizer, ::_RawLegacyNumberOfSweeps, value)
MOI.set(sampler, NumberOfSweeps(), value)
return nothing
end

function MOI.set(sampler::Optimizer, ::_RawLegacyNumberOfReplicas, value)
MOI.set(sampler, NumberOfReplicas(), value)
return nothing
end

function MOI.set(sampler::Optimizer, ::_RawLegacyNumberOfReads, value)
MOI.set(sampler, NumberOfReads(), value)
return nothing
end

function MOI.set(sampler::Optimizer, ::_RawMinimumTemperature, value)
MOI.set(sampler, MinimumTemperature(), value)
return nothing
end

function MOI.set(sampler::Optimizer, ::_RawMaximumTemperature, value)
MOI.set(sampler, MaximumTemperature(), value)
return nothing
end

MOI.supports(::Optimizer, ::_RawLegacyNumberOfSweeps) = true
MOI.supports(::Optimizer, ::_RawLegacyNumberOfReplicas) = true
MOI.supports(::Optimizer, ::_RawLegacyNumberOfReads) = true
MOI.supports(::Optimizer, ::_RawMinimumTemperature) = true
MOI.supports(::Optimizer, ::_RawMaximumTemperature) = true

QUBODrivers.honors_final_reads(::Type{<:Optimizer}) = true
QUBODrivers.enforces_time_limit(::Type{<:Optimizer}) = false

function _float_type(::Type{T})::String where {T<:AbstractFloat}
if T === Float16
return "float16"
Expand All @@ -46,6 +111,46 @@ function _float_type(::Type{T})::String where {T<:AbstractFloat}
end
end

function _check_nonnegative_integer(name::String, value::Integer)
value >= 0 || error("Value for '$name' must be a non-negative integer")

return nothing
end

function _seed_backend!(::Nothing)
return nothing
end

function _seed_backend!(seed::Integer)
np.random.seed(seed)
seed_numba_rng(seed)

return nothing
end

function _pysa_spin_sample(state, h, J, α, β, ::Type{T}) where {T}
ψ = -round.(Int, pyconvert.(T, [var for var in state]))
λ = QUBOTools.value(ψ, h, J, α, β)

return QUBOTools.Sample{T,Int}(ψ, λ)
end

function _final_samples(result, final_num_reads::Integer, h, J, α, β, ::Type{T}) where {T}
best_states = result["best_state"].values
backend_num_reads = length(best_states)
samples = Vector{QUBOTools.Sample{T,Int}}(undef, backend_num_reads)

for i = 1:backend_num_reads
# NOTE: Python is 0-indexed
# NOTE: sign has to be inverted, as mentioned before
samples[i] = _pysa_spin_sample(best_states[i-1], h, J, α, β, T)
end

sort!(samples; by = QUBOTools.value)

return samples[1:final_num_reads]
end

function QUBODrivers.sample(sampler::Optimizer{T}) where {T}
n, h, J, α, β = QUBOTools.ising(sampler, :dense; sense = :min)

Expand All @@ -60,42 +165,76 @@ function QUBODrivers.sample(sampler::Optimizer{T}) where {T}
float_type = _float_type(T),
)

num_sweeps = MOI.get(sampler, PySA.NumberOfSweeps())
num_replicas = MOI.get(sampler, PySA.NumberOfReplicas())
num_reads = MOI.get(sampler, PySA.NumberOfReads())
final_num_reads = MOI.get(sampler, QUBODrivers.FinalNumberOfReads())
seed = MOI.get(sampler, QUBODrivers.RandomSeed())
requested_parallel = MOI.get(sampler, PySA.Parallel())
backend_parallel = requested_parallel && isnothing(seed)
backend_num_reads = max(num_reads, final_num_reads)

_check_nonnegative_integer("NumberOfSweeps", num_sweeps)
_check_nonnegative_integer("NumberOfReplicas", num_replicas)
_check_nonnegative_integer("NumberOfReads", num_reads)
_check_nonnegative_integer("FinalNumberOfReads", final_num_reads)
num_replicas > 0 || error("Value for 'NumberOfReplicas' must be positive")

result = @timed begin
if backend_num_reads == 0
nothing
else
_seed_backend!(seed)
solver.metropolis_update(
num_sweeps = num_sweeps,
num_reads = backend_num_reads,
num_replicas = num_replicas,
update_strategy = MOI.get(sampler, PySA.UpdateStrategy()),
min_temp = MOI.get(sampler, PySA.MinimumTemperature()),
max_temp = MOI.get(sampler, PySA.MaximumTemperature()),
initialize_strategy = MOI.get(sampler, PySA.InitializeStrategy()),
recompute_energy = MOI.get(sampler, PySA.RecomputeEnergy()),
sort_output_temps = MOI.get(sampler, PySA.SortOutputTemps()),
parallel = backend_parallel,
verbose = !MOI.get(sampler, MOI.Silent()),
)
end
end

result = @timed solver.metropolis_update(
num_sweeps = MOI.get(sampler, PySA.NumberOfSweeps()),
num_reads = MOI.get(sampler, PySA.NumberOfReads()),
num_replicas = num_replicas,
update_strategy = MOI.get(sampler, PySA.UpdateStrategy()),
min_temp = MOI.get(sampler, PySA.MinimumTemperature()),
max_temp = MOI.get(sampler, PySA.MaximumTemperature()),
initialize_strategy = MOI.get(sampler, PySA.InitializeStrategy()),
recompute_energy = MOI.get(sampler, PySA.RecomputeEnergy()),
sort_output_temps = MOI.get(sampler, PySA.SortOutputTemps()),
parallel = MOI.get(sampler, PySA.Parallel()), # True by default
verbose = !MOI.get(sampler, MOI.Silent()),
)

samples = Vector{QUBOTools.Sample{T,Int}}(undef, num_replicas)

for Ψ in result.value["states"].values
for i = 1:num_replicas
# NOTE: Python is 0-indexed
# NOTE: sign has to be inverted, as mentioned before
ψ = -round.(Int, pyconvert.(T, [var for var in Ψ[i-1]]))

# Compute value instead of retrieving it, to avoid precision errors
λ = QUBOTools.value(ψ, h, J, α, β)
samples = if backend_num_reads == 0
QUBOTools.Sample{T,Int}[]
else
_final_samples(result.value, final_num_reads, h, J, α, β, T)
end

samples[i] = QUBOTools.Sample{T,Int}(ψ, λ)
end
seeds = if isnothing(seed)
Dict{String,Any}()
else
Dict{String,Any}("numpy" => seed, "numba" => seed)
end

metadata = Dict{String,Any}(
"origin" => "PySA",
"time" => Dict{String,Any}(
"effective" => result.time
),
metadata = QUBODrivers._sampler_metadata(
origin = "PySA.jl",
algorithm_name = "simulated_annealing",
backend_name = "pysa",
backend_version = MOI.get(sampler, MOI.SolverVersion()),
execution_mode = backend_parallel ? "parallel" : "sequential",
optimizer_iterations = num_sweeps,
optimizer_evaluations = backend_num_reads * num_replicas,
number_of_reads = backend_num_reads,
final_number_of_reads = final_num_reads,
seeds = seeds,
status = "locally_solved",
termination_status = MOI.LOCALLY_SOLVED,
)
metadata["time"] = Dict{String,Any}("effective" => result.time)
metadata["parameters"] = Dict{String,Any}(
"num_sweeps" => num_sweeps,
"num_replicas" => num_replicas,
"num_reads" => num_reads,
"final_num_reads" => final_num_reads,
"requested_parallel" => requested_parallel,
"parallel" => backend_parallel,
)

return QUBOTools.SampleSet{T}(samples, metadata; domain = :spin, sense = :min)
Expand Down
Loading
Loading