diff --git a/qubosolver/solving/classical/simulated_annealing.py b/qubosolver/solving/classical/simulated_annealing.py index 3dbc1cf9..847880ba 100644 --- a/qubosolver/solving/classical/simulated_annealing.py +++ b/qubosolver/solving/classical/simulated_annealing.py @@ -28,6 +28,8 @@ vectori, ) +from qubosolver.utils._costs import batched_quadratic_cost, _flip_deltas + logger = logging.getLogger(__name__) @@ -65,6 +67,12 @@ def _shrink(visited_solutions: dict[bytes, _Data], *, top_k: int) -> None: # picked, instead of each `def` capturing its own independent instance. _default_rng = torch_rng() +# How often the vectorized runner recomputes `energy` and `QX` exactly instead +# of accumulating them incrementally. Small enough that rounding error stays +# well below the tolerance of `Solution.check_consistency`, large enough that +# the extra matmul stays negligible against the per-iteration cost. +_REFRESH_EVERY = 128 + @overload def solve( @@ -80,6 +88,7 @@ def solve( time_limit: float = float("inf"), rng: torch.Generator = _default_rng, stats: Literal["per_run", "full"] = "per_run", + vectorized: bool = True, ) -> Solution: ... @@ -97,6 +106,7 @@ def solve( time_limit: float = float("inf"), rng: torch.Generator = _default_rng, stats: Literal["per_run", "full"] = "per_run", + vectorized: bool = True, ) -> list[Solution]: ... @@ -114,6 +124,7 @@ def solve( time_limit: float = float("inf"), rng: torch.Generator = _default_rng, stats: Literal["per_run", "full"] = "per_run", + vectorized: bool = True, ) -> Solution | list[Solution]: """Run Simulated Annealing on a QUBO instance from each of a batch of starting points. @@ -169,10 +180,15 @@ def solve( time_limit: Wall-clock budget in seconds. The algorithm stops early when either `max_iter` steps or the time limit is reached, whichever comes first. Defaults to ``float("inf")`` (no limit). + With ``vectorized=True`` this is a single budget for the whole + batch of starts, which all stop at the same iteration; with + ``vectorized=False`` each start gets its own budget. rng: PyTorch random number generator used for bit selection and acceptance sampling. Defaults to a module-level generator created once at import time; pass an explicit generator - for reproducibility across calls. + for reproducibility across calls. Note that `vectorized` changes + the order in which draws are consumed, so a given seed produces + the same result only for a fixed value of `vectorized`. stats: When ``"per_run"`` (default), each run's retained bitstrings are counted as ``1`` instead of how many iterations were spent at each one, before any merging. This is mainly meant for @@ -184,6 +200,13 @@ def solve( generally neither ``1`` nor uniform. When ``"full"``, counts instead reflect how many iterations were spent at each bitstring. + vectorized: When ``True`` (default), step every start forward together + so each iteration costs a handful of batched tensor operations + regardless of how many starts there are -- markedly faster for + large batches. When ``False``, anneal the starts one at a time; + this is the reference implementation, kept for comparison. The two + run the same algorithm but differ in `time_limit` scope and in RNG + draw order (see `time_limit` and `rng`). Returns: When ``merge=True``, a single [`Solution`][] merging every start's @@ -211,7 +234,6 @@ def solve( ) n = instance.size - Q = instance.matrix # determine cooling rate alpha if max_iter <= 1: @@ -223,10 +245,73 @@ def solve( else: alpha = (final_temp / initial_temp) ** (1.0 / (max_iter - 1)) - solutions: list[Solution] = [] if isinstance(starts, int): starts = bitstrings.rand(starts, n, rng=rng) + if starts.shape[0] == 0: + return Solution() if merge else [] + + runner = _run_vectorized if vectorized else _run_sequential + solutions = runner( + instance, + starts, + top_k=top_k, + max_iter=max_iter, + initial_temp=initial_temp, + alpha=alpha, + time_limit=time_limit, + rng=rng, + stats=stats, + ) + + if merge: + return Solution.concat(solutions).deduplicate() + + return solutions + + +def _run_sequential( + instance: Instance, + starts: Bitstrings, + *, + top_k: int, + max_iter: int, + initial_temp: float, + alpha: float, + time_limit: float, + rng: torch.Generator, + stats: Literal["per_run", "full"], +) -> list[Solution]: + """Anneal each start in turn, one scalar bit-flip proposal at a time. + + The original, run-at-a-time implementation, kept as the reference + behaviour that [`_run_vectorized`][] is checked against. It is + `O(len(starts) * max_iter)` in Python-level torch calls, so prefer the + vectorized path for large batches. + + Unlike the vectorized path, `time_limit` here is consumed per run: each + start gets its own fresh deadline. + + Args: + instance: The QUBO instance to solve; its coefficient matrix must + already be symmetric. + starts: Batch of initial bitstrings of shape ``(k, n)``. + top_k: Maximum number of unique best solutions to keep per run. + max_iter: Number of bit-flip proposals per run. + initial_temp: Starting temperature. + alpha: Geometric cooling factor applied at each step. + time_limit: Wall-clock budget in seconds, per run. + rng: Generator used for bit selection and acceptance sampling. + stats: See [`solve`][]. + + Returns: + One [`Solution`][] per row of `starts`, in the same order, each + sorted by ascending cost with probabilities computed. + """ + Q = instance.matrix + n = Q.shape[0] + solutions: list[Solution] = [] + for b in starts: bits: Bitstring = b.detach().clone() @@ -239,11 +324,19 @@ def solve( visited_solutions[_to_key(bits)] = _Data(energy, 1) deadline = time.perf_counter() + time_limit + visits = 1 for _ in range(max_iter): if time.perf_counter() >= deadline: break + # See the matching comment in `_run_vectorized`: `energy` and `Qx` + # are accumulated incrementally, so periodically recomputing them + # exactly from `bits` keeps rounding error from growing unbounded. + if visits % _REFRESH_EVERY == 0: + Qx = Q @ bits.to(Q) + energy = float(bits.to(Q).dot(Qx)) + i = int(torch.randint(0, n, (1,), generator=rng).item()) xi = int(bits[i].item()) @@ -266,6 +359,7 @@ def solve( key = _to_key(bits) sol = visited_solutions.setdefault(key, _Data(energy, 0)) sol.count += 1 + visits += 1 # Most inserts are one-off bitstrings that will never make the # top_k cut, so the dict keeps growing between shrinks regardless @@ -293,9 +387,147 @@ def solve( counts=counts, ) - solutions.append(solution._sort_by_cost()._compute_probabilities()) + # `costs` was accumulated incrementally, drifting from the true x^T Q x + # by up to _REFRESH_EVERY steps of rounding error; recompute it exactly + # now that the hot loop is done. + solutions.append(solution._update(instance)) - if merge: - return Solution.concat(solutions).deduplicate() + return solutions + + +def _run_vectorized( + instance: Instance, + starts: Bitstrings, + *, + top_k: int, + max_iter: int, + initial_temp: float, + alpha: float, + time_limit: float, + rng: torch.Generator, + stats: Literal["per_run", "full"], +) -> list[Solution]: + """Anneal every start simultaneously, one batched bit-flip proposal per step. + + All runs are stepped forward together: each iteration proposes one flip per + run, evaluates every candidate delta in a single matmul, and accepts or + rejects per run via a boolean mask. The number of Python-level torch calls + is therefore proportional to `max_iter` alone rather than to + ``len(starts) * max_iter``, so the cost is nearly flat in the number of + starts. The runs remain statistically independent -- only their bookkeeping + is shared. + + Two differences from [`_run_sequential`][] follow from batching: + + - `time_limit` is a single budget for the whole batch, checked once per + iteration, so all runs stop at the same iteration. + - Random draws are batched across runs, so a given seed does not reproduce + the sequential path's draw order. + + Args: + instance: See [`_run_sequential`][]. + starts: See [`_run_sequential`][]. + top_k: See [`_run_sequential`][]. + max_iter: See [`_run_sequential`][]. + initial_temp: See [`_run_sequential`][]. + alpha: See [`_run_sequential`][]. + time_limit: Wall-clock budget in seconds, for the whole batch. + rng: See [`_run_sequential`][]. + stats: See [`solve`][]. + + Returns: + One [`Solution`][] per row of `starts`, in the same order, each + sorted by ascending cost with probabilities computed. + """ + Q = instance.matrix + n = Q.shape[0] + n_runs = starts.shape[0] + + X = starts.detach().clone().to(Q) + QX = X @ Q + energy = batched_quadratic_cost(X, Q) + + # Preallocated for the worst case (every iteration runs) and sliced down to + # however many actually did; each iteration writes the post-move state of + # every run at index `visits`, with index 0 holding the starting state. + visited_bits = torch.empty((max_iter + 1, n_runs, n), dtype=bitstring.dtype(), device=X.device) + visited_energy = torch.empty((max_iter + 1, n_runs), dtype=energy.dtype, device=X.device) + visited_bits[0] = X.to(bitstring.dtype()) + visited_energy[0] = energy + + temperature: float = initial_temp + rows = torch.arange(n_runs, device=X.device) + deadline = time.perf_counter() + time_limit + visits = 1 + + for _ in range(max_iter): + if time.perf_counter() >= deadline: + break + + # `energy` and `QX` are accumulated incrementally, so each step adds a + # rounding error that leaves the reported costs a few ULPs off the true + # x^T Q x. Recomputing them exactly every _REFRESH_EVERY steps keeps + # that error from becoming visible, at the cost of one extra matmul + # amortized over many iterations. `X` itself never needs this: each + # accepted flip moves it by an exact +1/-1, so it stays exact with no + # drift to correct. + if visits % _REFRESH_EVERY == 0: + QX = X @ Q + energy = batched_quadratic_cost(X, Q) + + idx = torch.randint(0, n, (n_runs,), generator=rng, device=X.device) + dE = _flip_deltas(Q, X, QX).gather(1, idx.unsqueeze(1)).squeeze(1) + + accept = (dE <= 0.0) | ( + torch.rand(n_runs, generator=rng, device=X.device, dtype=X.dtype) + < torch.exp(-dE / temperature) + ) + + # Flipping x -> 1 - x moves the bit by +1 or -1; zeroing that step on + # the rejected runs leaves them untouched without branching per run. + xi = X[rows, idx] + step = (1.0 - 2.0 * xi) * accept.to(X.dtype) + X[rows, idx] = xi + step + energy = energy + step.abs() * dE + QX = QX + step.unsqueeze(1) * Q[idx, :] + + visited_bits[visits] = X.to(bitstring.dtype()) + visited_energy[visits] = energy + visits += 1 + + temperature *= alpha + if temperature < 1e-12: + temperature = 1e-12 + + # (visits, n_runs, n) and (visits, n_runs): one entry per recorded state. + visited_bits = visited_bits[:visits] + visited_energy = visited_energy[:visits] + + solutions: list[Solution] = [] + for r in range(n_runs): + # Every recorded state starts as its own candidate counting a single + # visit; `deduplicate` then collapses repeats of the same bitstring, + # summing those visits into its count and keeping its energy. Since it + # leaves the result sorted by ascending cost, `truncate` reduces it to + # the top_k lowest-energy ones. This loop runs once per run rather than + # once per iteration, so it stays off the hot path. + solution = Solution( + bitstrings=bitstrings.as_tensor(visited_bits[:, r, :]), + costs=visited_energy[:, r], + counts=vectori.zeros(visits).fill_(1), + probabilities=vector.zeros(visits).fill_(1.0 / visits), + ) + # `costs` was accumulated incrementally, drifting from the true x^T Q x + # by up to _REFRESH_EVERY steps of rounding error. `deduplicate` picks + # the row to keep per bitstring based on that drifted cost, so skip its + # own recompute (`update=False`) and instead recompute exactly via + # `_update` right after, before `truncate` selects on it. + solution.deduplicate(update=False)._update(instance).truncate(top_k) + + if stats == "per_run": + solution.counts = solution.counts.clone().fill_(1) + solution._compute_probabilities() + + solutions.append(solution) return solutions diff --git a/tests/solvers/test_simulated_annealing.py b/tests/solvers/test_simulated_annealing.py index fc309205..b385d950 100644 --- a/tests/solvers/test_simulated_annealing.py +++ b/tests/solvers/test_simulated_annealing.py @@ -9,6 +9,7 @@ from typing_extensions import assert_type, get_overloads from qubosolver import ( + Dataset, Instance, Solution, bitstring, @@ -53,6 +54,21 @@ instances = [instance_symmetric, instance_small] instance_ids = ["6var", "4var"] +vectorized_params = [ + pytest.param(True, id="vectorized"), + pytest.param(False, id="sequential"), +] + +# The non-vectorized path does one Python-level torch call per (start, iteration) +# pair, so its wall-clock cost scales with n_starts * max_iter regardless of +# instance size; skip it once that product gets large enough to slow the suite. +_MAX_SEQUENTIAL_WORK = 2000 + + +def skip_if_sequential_too_slow(vectorized: bool, *, n_starts: int, max_iter: int) -> None: + if not vectorized and n_starts * max_iter > _MAX_SEQUENTIAL_WORK: + pytest.skip("n_starts * max_iter too large for the non-vectorized path") + def test_to_key_from_key_round_trip() -> None: """_from_key must reconstruct the exact bitstring given to _to_key.""" @@ -128,8 +144,10 @@ def test_shrink_noop_when_already_within_top_k() -> None: @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_costs_match_bitstrings(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_costs_match_bitstrings(instance: Instance, vectorized: bool) -> None: """Every reported cost must correspond to x^T Q x of its own bitstring.""" + skip_if_sequential_too_slow(vectorized, n_starts=1, max_iter=3000) start = bitstrings.zeros(1, instance.size) rng = torch_rng(0) @@ -141,6 +159,7 @@ def test_simulated_annealing_costs_match_bitstrings(instance: Instance) -> None: initial_temp=4.0, final_temp=0.05, rng=rng, + vectorized=vectorized, ) true_solution = copy.deepcopy(solution)._compute_costs(instance.matrix) @@ -151,11 +170,34 @@ def test_simulated_annealing_costs_match_bitstrings(instance: Instance) -> None: ) +@pytest.mark.priority(5) +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_energy_does_not_drift_on_large_batch(vectorized: bool) -> None: + """On a large enough batch/instance, accumulated rounding error in the + incrementally tracked `energy` must not make reported costs drift from the + true `x^T Q x`, for either path: both `_run_sequential` and + `_run_vectorized` periodically recompute it exactly (see `_REFRESH_EVERY`), + and recompute it once more from the final bitstrings before returning.""" + rng = torch_rng(64548) + dataset = Dataset.from_random(1, 100, rng=rng) + instance, _ = dataset[0] + starts = bitstrings.rand(400, instance.size, rng=rng) + + solution = solving.simulated_annealing.solve( + instance, starts=starts, max_iter=1000, vectorized=vectorized, rng=rng + ) + check.is_true(solution.check_consistency(instance=instance, throw=False)) + + @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_solution_is_internally_consistent(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_solution_is_internally_consistent( + instance: Instance, vectorized: bool +) -> None: """The returned Solution must pass the full consistency check (shapes, costs, sortedness, no duplicate bitstrings, positive integer counts, probabilities matching normalised counts).""" + skip_if_sequential_too_slow(vectorized, n_starts=1, max_iter=3000) start = bitstrings.zeros(1, instance.size) rng = torch_rng(0) @@ -167,13 +209,15 @@ def test_simulated_annealing_solution_is_internally_consistent(instance: Instanc initial_temp=4.0, final_temp=0.05, rng=rng, + vectorized=vectorized, ) check.is_true(solution.check_consistency(instance=instance, throw=True)) @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_counts_sum_to_visits(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_counts_sum_to_visits(instance: Instance, vectorized: bool) -> None: """Counts must be strictly positive integers, and their total must be at least the number of returned bitstrings.""" start = bitstrings.zeros(1, instance.size) @@ -187,6 +231,7 @@ def test_simulated_annealing_counts_sum_to_visits(instance: Instance) -> None: initial_temp=4.0, final_temp=0.05, rng=rng, + vectorized=vectorized, ) check.is_true(torch.all(solution.counts > 0).item()) @@ -194,7 +239,8 @@ def test_simulated_annealing_counts_sum_to_visits(instance: Instance) -> None: @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_respects_top_k(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_respects_top_k(instance: Instance, vectorized: bool) -> None: """The number of returned solutions never exceeds top_k.""" start = bitstrings.zeros(1, instance.size) rng = torch_rng(0) @@ -207,13 +253,17 @@ def test_simulated_annealing_respects_top_k(instance: Instance) -> None: initial_temp=4.0, final_temp=0.05, rng=rng, + vectorized=vectorized, ) check.is_in(len(solution), [1, 2]) @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_deterministic_with_seeded_rng(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_deterministic_with_seeded_rng( + instance: Instance, vectorized: bool +) -> None: """Two runs with the same seed must produce identical solutions.""" start = bitstrings.zeros(1, instance.size) @@ -225,6 +275,7 @@ def test_simulated_annealing_deterministic_with_seeded_rng(instance: Instance) - initial_temp=4.0, final_temp=0.05, rng=torch_rng(565111), + vectorized=vectorized, ) solution_b = solving.simulated_annealing.solve( instance, @@ -234,6 +285,7 @@ def test_simulated_annealing_deterministic_with_seeded_rng(instance: Instance) - initial_temp=4.0, final_temp=0.05, rng=torch_rng(565111), + vectorized=vectorized, ) torch.testing.assert_close(solution_a.bitstrings, solution_b.bitstrings) @@ -245,7 +297,10 @@ def test_simulated_annealing_deterministic_with_seeded_rng(instance: Instance) - @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_zero_max_iter_returns_start(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_zero_max_iter_returns_start( + instance: Instance, vectorized: bool +) -> None: """With max_iter=0, only the starting bitstring is returned.""" start = bitstrings.zeros(1, instance.size) rng = torch_rng(0) @@ -258,6 +313,7 @@ def test_simulated_annealing_zero_max_iter_returns_start(instance: Instance) -> initial_temp=4.0, final_temp=0.05, rng=rng, + vectorized=vectorized, ) check.equal(len(solution), 1) @@ -266,7 +322,10 @@ def test_simulated_annealing_zero_max_iter_returns_start(instance: Instance) -> @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_zero_time_limit_returns_start(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_zero_time_limit_returns_start( + instance: Instance, vectorized: bool +) -> None: """An exhausted time budget stops the loop before any iteration runs.""" start = bitstrings.zeros(1, instance.size) rng = torch_rng(0) @@ -280,6 +339,7 @@ def test_simulated_annealing_zero_time_limit_returns_start(instance: Instance) - final_temp=0.05, time_limit=0.0, rng=rng, + vectorized=vectorized, ) check.equal(len(solution), 1) @@ -287,7 +347,10 @@ def test_simulated_annealing_zero_time_limit_returns_start(instance: Instance) - @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_explicit_cooling_rate_used(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_explicit_cooling_rate_used( + instance: Instance, vectorized: bool +) -> None: """When cooling_rate is provided, final_temp is ignored and no error is raised even if final_temp is invalid (<= 0).""" start = bitstrings.zeros(1, instance.size) @@ -302,6 +365,7 @@ def test_simulated_annealing_explicit_cooling_rate_used(instance: Instance) -> N final_temp=-1.0, cooling_rate=0.9, rng=rng, + vectorized=vectorized, ) check.is_true(solution.check_consistency(instance=instance, throw=True)) @@ -332,8 +396,9 @@ def test_simulated_annealing_raises_on_invalid_arguments(kwargs: dict, match: st @pytest.mark.parametrize("instance", instances, ids=instance_ids) +@pytest.mark.parametrize("vectorized", vectorized_params) def test_simulated_annealing_merge_false_returns_one_solution_per_start( - instance: Instance, + instance: Instance, vectorized: bool ) -> None: """With merge=False, one Solution must be returned per row of `start`, in the same order, none of them merged with the others.""" @@ -349,6 +414,7 @@ def test_simulated_annealing_merge_false_returns_one_solution_per_start( initial_temp=4.0, final_temp=0.05, rng=rng, + vectorized=vectorized, ) check.equal(len(solutions), 3) @@ -357,8 +423,9 @@ def test_simulated_annealing_merge_false_returns_one_solution_per_start( @pytest.mark.parametrize("instance", instances, ids=instance_ids) +@pytest.mark.parametrize("vectorized", vectorized_params) def test_simulated_annealing_merge_true_matches_manual_concat_and_deduplicate( - instance: Instance, + instance: Instance, vectorized: bool ) -> None: """merge=True (the default) must be equivalent to merging the merge=False per-start results via Solution.concat(...).deduplicate(), as documented @@ -374,6 +441,7 @@ def test_simulated_annealing_merge_true_matches_manual_concat_and_deduplicate( initial_temp=4.0, final_temp=0.05, rng=torch_rng(7874), + vectorized=vectorized, ) solutions = solving.simulated_annealing.solve( instance, @@ -384,6 +452,7 @@ def test_simulated_annealing_merge_true_matches_manual_concat_and_deduplicate( initial_temp=4.0, final_temp=0.05, rng=torch_rng(7874), + vectorized=vectorized, ) manually_merged = Solution.concat(solutions).deduplicate() @@ -393,7 +462,8 @@ def test_simulated_annealing_merge_true_matches_manual_concat_and_deduplicate( torch.testing.assert_close(merged_solution.counts, manually_merged.counts) -def test_simulated_annealing_int_starts_generates_that_many_random_runs() -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_int_starts_generates_that_many_random_runs(vectorized: bool) -> None: """Passing an int for `starts` must generate that many uniformly random starting bitstrings from `rng`, giving the same result as pre-generating them with bitstrings.rand from the same rng and passing them explicitly.""" @@ -407,6 +477,7 @@ def test_simulated_annealing_int_starts_generates_that_many_random_runs() -> Non top_k=2, max_iter=50, rng=rng_int, + vectorized=vectorized, ) rng_explicit = torch_rng(5821) @@ -418,6 +489,7 @@ def test_simulated_annealing_int_starts_generates_that_many_random_runs() -> Non top_k=2, max_iter=50, rng=rng_explicit, + vectorized=vectorized, ) check.equal(len(solution), n_starts) @@ -426,11 +498,17 @@ def test_simulated_annealing_int_starts_generates_that_many_random_runs() -> Non torch.testing.assert_close(actual.costs, exp.costs, atol=0.0, rtol=0.0) -def test_simulated_annealing_default_starts_is_one_random_start() -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_default_starts_is_one_random_start(vectorized: bool) -> None: """Omitting `starts` must default to a single uniformly random start, producing exactly one run's worth of results.""" solution = solving.simulated_annealing.solve( - instance_symmetric, merge=False, top_k=1, max_iter=50, rng=torch_rng(0) + instance_symmetric, + merge=False, + top_k=1, + max_iter=50, + rng=torch_rng(0), + vectorized=vectorized, ) check.equal(len(solution), 1) @@ -460,8 +538,9 @@ def test_simulated_annealing_empty_start_merge_true_returns_empty_solution() -> @pytest.mark.parametrize("instance", instances, ids=instance_ids) +@pytest.mark.parametrize("vectorized", vectorized_params) def test_simulated_annealing_stats_per_run_sets_single_run_counts_to_one( - instance: Instance, + instance: Instance, vectorized: bool ) -> None: """With a single run (one start), stats='per_run' must set every returned bitstring's count to 1, regardless of how many iterations were @@ -481,16 +560,21 @@ def test_simulated_annealing_stats_per_run_sets_single_run_counts_to_one( final_temp=0.05, rng=rng, stats="per_run", + vectorized=vectorized, ) expected_counts = vectori.zeros(len(solution)).fill_(1) torch.testing.assert_close(solution.counts, expected_counts) -def test_simulated_annealing_stats_per_run_merged_counts_reflect_run_agreement() -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_stats_per_run_merged_counts_reflect_run_agreement( + vectorized: bool, +) -> None: """With stats='per_run', top_k=1, and merge=True (default), each run contributes a single bitstring with count 1; after merging, a bitstring's count is the number of runs that converged on it -- neither always 1 nor uniform across bitstrings.""" + skip_if_sequential_too_slow(vectorized, n_starts=8, max_iter=300) start = bitstrings.rand(8, instance_symmetric.size, rng=torch_rng(11)) solution = solving.simulated_annealing.solve( @@ -502,13 +586,17 @@ def test_simulated_annealing_stats_per_run_merged_counts_reflect_run_agreement() final_temp=0.05, rng=torch_rng(0), stats="per_run", + vectorized=vectorized, ) check.equal(solution.counts.sum().item(), 8) check.is_true(torch.all(solution.counts >= 1).item()) -def test_simulated_annealing_stats_per_run_top_k_one_merge_true_matches_manual_equivalent() -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_stats_per_run_top_k_one_merge_true_matches_manual_equivalent( + vectorized: bool, +) -> None: """merge=True, top_k=1, stats='per_run' must be equivalent to running with merge=False, top_k>1, stats='full' (the default), then per start keeping only the best bitstring (truncate(1) -- each per-start Solution @@ -520,6 +608,7 @@ def test_simulated_annealing_stats_per_run_top_k_one_merge_true_matches_manual_e most other optimization libraries return by default, while still running with stats='full' to keep the complete per-run results available if needed.""" + skip_if_sequential_too_slow(vectorized, n_starts=8, max_iter=300) start = bitstrings.rand(8, instance_symmetric.size, rng=torch_rng(1350)) per_run_solution = solving.simulated_annealing.solve( @@ -529,6 +618,7 @@ def test_simulated_annealing_stats_per_run_top_k_one_merge_true_matches_manual_e initial_temp=4.0, final_temp=0.05, rng=torch_rng(0), + vectorized=vectorized, ) solutions = solving.simulated_annealing.solve( @@ -541,6 +631,7 @@ def test_simulated_annealing_stats_per_run_top_k_one_merge_true_matches_manual_e final_temp=0.05, rng=torch_rng(0), stats="full", + vectorized=vectorized, ) manually_equivalent = Solution.concat( [solution.truncate(1) for solution in solutions], unit_counts=True @@ -557,7 +648,8 @@ def test_simulated_annealing_stats_per_run_top_k_one_merge_true_matches_manual_e @pytest.mark.parametrize("instance", instances, ids=instance_ids) -def test_simulated_annealing_stats_per_run_is_default(instance: Instance) -> None: +@pytest.mark.parametrize("vectorized", vectorized_params) +def test_simulated_annealing_stats_per_run_is_default(instance: Instance, vectorized: bool) -> None: """Omitting stats must be equivalent to passing stats='per_run' explicitly.""" start = bitstrings.zeros(1, instance.size) @@ -569,6 +661,7 @@ def test_simulated_annealing_stats_per_run_is_default(instance: Instance) -> Non initial_temp=4.0, final_temp=0.05, rng=torch_rng(23), + vectorized=vectorized, ) explicit_per_run_solution = solving.simulated_annealing.solve( instance, @@ -579,6 +672,7 @@ def test_simulated_annealing_stats_per_run_is_default(instance: Instance) -> Non final_temp=0.05, rng=torch_rng(23), stats="per_run", + vectorized=vectorized, ) torch.testing.assert_close(default_solution.counts, explicit_per_run_solution.counts)