1use scirs2_core::ndarray::{Array, ArrayD, Ix2};
85use scirs2_core::random::prelude::*;
86use scirs2_core::random::rngs::StdRng;
87use std::collections::HashMap;
88
89use super::energy::{compute_influence_simd, energy_full_simd, update_influence_simd};
90use super::{SampleResult, Sampler, SamplerError, SamplerResult};
91
92#[derive(Debug, Clone)]
94pub struct TabuParams {
95 pub tenure: usize,
97 pub max_iter: usize,
99 pub restart_threshold: usize,
101 pub aspiration: bool,
103}
104
105impl Default for TabuParams {
106 fn default() -> Self {
107 Self {
108 tenure: 10,
109 max_iter: 1000,
110 restart_threshold: 200,
111 aspiration: true,
112 }
113 }
114}
115
116#[derive(Debug, Clone)]
147pub struct TabuSampler {
148 pub seed: Option<u64>,
150 pub params: TabuParams,
152}
153
154impl TabuSampler {
155 #[must_use]
157 pub fn new() -> Self {
158 Self {
159 seed: None,
160 params: TabuParams::default(),
161 }
162 }
163
164 #[must_use]
166 pub fn with_seed(mut self, seed: u64) -> Self {
167 self.seed = Some(seed);
168 self
169 }
170
171 #[must_use]
173 pub fn with_tenure(mut self, tenure: usize) -> Self {
174 self.params.tenure = tenure;
175 self
176 }
177
178 #[must_use]
180 pub fn with_max_iter(mut self, max_iter: usize) -> Self {
181 self.params.max_iter = max_iter;
182 self
183 }
184
185 #[must_use]
187 pub fn with_restart_threshold(mut self, restart_threshold: usize) -> Self {
188 self.params.restart_threshold = restart_threshold;
189 self
190 }
191
192 #[must_use]
194 pub fn with_aspiration(mut self, aspiration: bool) -> Self {
195 self.params.aspiration = aspiration;
196 self
197 }
198
199 #[inline]
203 fn compute_qubo_energy(q_matrix: &[f64], state: &[bool], n: usize) -> f64 {
204 energy_full_simd(state, q_matrix, n)
205 }
206
207 #[inline]
213 fn compute_influence_vector(q_matrix: &[f64], state: &[bool], n: usize) -> Vec<f64> {
214 compute_influence_simd(state, q_matrix, n)
215 }
216
217 #[inline]
223 fn update_influence_vector(
224 g: &mut [f64],
225 q_matrix: &[f64],
226 k: usize,
227 flipped_to: bool,
228 n: usize,
229 ) {
230 update_influence_simd(g, q_matrix, n, k, flipped_to);
231 }
232
233 fn run_single_qubo(
235 &self,
236 q_matrix: &[f64],
237 n_vars: usize,
238 rng: &mut StdRng,
239 ) -> (Vec<bool>, f64) {
240 if n_vars == 0 {
241 return (vec![], 0.0);
242 }
243
244 let tenure = self.params.tenure;
245 let max_iter = self.params.max_iter;
246 let restart_threshold = self.params.restart_threshold;
247 let aspiration = self.params.aspiration;
248
249 let mut state: Vec<bool> = (0..n_vars).map(|_| rng.random_bool(0.5)).collect();
251 let mut energy = Self::compute_qubo_energy(q_matrix, &state, n_vars);
252 let mut g = Self::compute_influence_vector(q_matrix, &state, n_vars);
253
254 let mut best_state = state.clone();
256 let mut best_energy = energy;
257
258 let mut tabu_until = vec![0usize; n_vars];
260 let mut non_improving_count = 0usize;
261 let mut iter = 0usize;
262
263 while iter < max_iter {
264 let mut best_flip_idx = None;
266 let mut best_flip_delta = f64::INFINITY;
267
268 for i in 0..n_vars {
269 let delta_e = (1.0 - 2.0 * if state[i] { 1.0 } else { 0.0 }) * g[i];
270 let is_tabu = tabu_until[i] > iter;
271
272 if !is_tabu {
273 if delta_e < best_flip_delta {
274 best_flip_delta = delta_e;
275 best_flip_idx = Some(i);
276 }
277 } else if aspiration && (energy + delta_e < best_energy) {
278 if delta_e < best_flip_delta {
280 best_flip_delta = delta_e;
281 best_flip_idx = Some(i);
282 }
283 }
284 }
285
286 let flip_idx = match best_flip_idx {
287 Some(idx) => idx,
288 None => {
289 let mut forced_best = 0;
291 let mut forced_best_delta = f64::INFINITY;
292 for i in 0..n_vars {
293 let delta_e = (1.0 - 2.0 * if state[i] { 1.0 } else { 0.0 }) * g[i];
294 if delta_e < forced_best_delta {
295 forced_best_delta = delta_e;
296 forced_best = i;
297 }
298 }
299 forced_best
300 }
301 };
302
303 let new_x = !state[flip_idx];
305 energy += best_flip_delta;
306 Self::update_influence_vector(&mut g, q_matrix, flip_idx, new_x, n_vars);
307 state[flip_idx] = new_x;
308
309 tabu_until[flip_idx] = iter + tenure + 1;
311
312 if energy < best_energy - 1e-12 {
314 best_energy = energy;
315 best_state = state.clone();
316 non_improving_count = 0;
317 } else {
318 non_improving_count += 1;
319 }
320
321 if non_improving_count >= restart_threshold {
323 state = best_state.clone();
325 energy = best_energy;
326
327 let n_perturb = (n_vars as f64 / 10.0).ceil() as usize;
328 for _ in 0..n_perturb {
329 let idx = rng.random_range(0..n_vars);
330 state[idx] = !state[idx];
331 }
332 energy = Self::compute_qubo_energy(q_matrix, &state, n_vars);
333 g = Self::compute_influence_vector(q_matrix, &state, n_vars);
334 non_improving_count = 0;
335 tabu_until.fill(0);
337 }
338
339 iter += 1;
340 }
341
342 (best_state, best_energy)
343 }
344
345 fn evaluate_hobo_energy<D>(tensor: &Array<f64, D>, state: &[bool], n_vars: usize) -> f64
347 where
348 D: scirs2_core::ndarray::Dimension + 'static,
349 {
350 let dyn_tensor: ArrayD<f64> = tensor.to_owned().into_dyn();
352 Self::evaluate_hobo_energy_dyn(&dyn_tensor, state, n_vars)
353 }
354
355 fn evaluate_hobo_energy_dyn(tensor: &ArrayD<f64>, state: &[bool], n_vars: usize) -> f64 {
357 let mut energy = 0.0;
358 let shape = tensor.shape();
359 let ndim = shape.len();
360
361 if ndim == 2 {
362 let n0 = shape[0].min(n_vars);
363 let n1 = shape[1].min(n_vars);
364 for i in 0..n0 {
365 if !state[i] {
366 continue;
367 }
368 for j in 0..n1 {
369 if state[j] {
370 if let Some(&v) = tensor.get([i, j].as_slice()) {
371 energy += v;
372 }
373 }
374 }
375 }
376 } else if ndim == 3 {
377 let n0 = shape[0].min(n_vars);
378 let n1 = shape[1].min(n_vars);
379 let n2 = shape[2].min(n_vars);
380 for i in 0..n0 {
381 if !state[i] {
382 continue;
383 }
384 for j in 0..n1 {
385 if !state[j] {
386 continue;
387 }
388 for k in 0..n2 {
389 if state[k] {
390 if let Some(&v) = tensor.get([i, j, k].as_slice()) {
391 energy += v;
392 }
393 }
394 }
395 }
396 }
397 }
398 energy
400 }
401
402 fn run_single_hobo<D>(
404 &self,
405 tensor: &Array<f64, D>,
406 n_vars: usize,
407 rng: &mut StdRng,
408 ) -> (Vec<bool>, f64)
409 where
410 D: scirs2_core::ndarray::Dimension + 'static,
411 {
412 let dyn_tensor: ArrayD<f64> = tensor.to_owned().into_dyn();
414 self.run_single_hobo_dyn(&dyn_tensor, n_vars, rng)
415 }
416
417 fn run_single_hobo_dyn(
419 &self,
420 tensor: &ArrayD<f64>,
421 n_vars: usize,
422 rng: &mut StdRng,
423 ) -> (Vec<bool>, f64) {
424 if n_vars == 0 {
425 return (vec![], 0.0);
426 }
427
428 let tenure = self.params.tenure;
429 let max_iter = self.params.max_iter;
430 let restart_threshold = self.params.restart_threshold;
431 let aspiration = self.params.aspiration;
432
433 let mut state: Vec<bool> = (0..n_vars).map(|_| rng.random_bool(0.5)).collect();
434 let mut energy = Self::evaluate_hobo_energy_dyn(tensor, &state, n_vars);
435
436 let mut best_state = state.clone();
437 let mut best_energy = energy;
438
439 let mut tabu_until = vec![0usize; n_vars];
440 let mut non_improving_count = 0usize;
441
442 for iter in 0..max_iter {
443 let mut best_flip_idx = None;
444 let mut best_flip_delta = f64::INFINITY;
445
446 for i in 0..n_vars {
447 let is_tabu = tabu_until[i] > iter;
448 state[i] = !state[i];
449 let new_energy = Self::evaluate_hobo_energy_dyn(tensor, &state, n_vars);
450 let delta_e = new_energy - energy;
451 state[i] = !state[i]; if (!is_tabu || (aspiration && (energy + delta_e < best_energy)))
454 && delta_e < best_flip_delta
455 {
456 best_flip_delta = delta_e;
457 best_flip_idx = Some(i);
458 }
459 }
460
461 let flip_idx = match best_flip_idx {
462 Some(idx) => idx,
463 None => {
464 let mut forced_best = 0;
466 let mut forced_best_delta = f64::INFINITY;
467 for i in 0..n_vars {
468 state[i] = !state[i];
469 let new_energy = Self::evaluate_hobo_energy_dyn(tensor, &state, n_vars);
470 let delta_e = new_energy - energy;
471 state[i] = !state[i];
472 if delta_e < forced_best_delta {
473 forced_best_delta = delta_e;
474 forced_best = i;
475 }
476 }
477 forced_best
478 }
479 };
480
481 state[flip_idx] = !state[flip_idx];
482 energy += best_flip_delta;
483 tabu_until[flip_idx] = iter + tenure + 1;
484
485 if energy < best_energy - 1e-12 {
486 best_energy = energy;
487 best_state = state.clone();
488 non_improving_count = 0;
489 } else {
490 non_improving_count += 1;
491 }
492
493 if non_improving_count >= restart_threshold {
494 state = best_state.clone();
495 energy = best_energy;
496 let n_perturb = (n_vars as f64 / 10.0).ceil() as usize;
497 for _ in 0..n_perturb {
498 let idx = rng.random_range(0..n_vars);
499 state[idx] = !state[idx];
500 }
501 energy = Self::evaluate_hobo_energy_dyn(tensor, &state, n_vars);
502 non_improving_count = 0;
503 tabu_until.fill(0);
504 }
505 }
506
507 (best_state, best_energy)
508 }
509
510 fn run_generic<D>(
512 &self,
513 matrix_or_tensor: &Array<f64, D>,
514 var_map: &HashMap<String, usize>,
515 shots: usize,
516 ) -> SamplerResult<Vec<SampleResult>>
517 where
518 D: scirs2_core::ndarray::Dimension + 'static,
519 {
520 let shots = shots.max(1);
521 let n_vars = var_map.len();
522 if n_vars == 0 {
523 return Err(SamplerError::InvalidParameter(
524 "Variable map is empty".to_string(),
525 ));
526 }
527
528 let idx_to_var: HashMap<usize, String> = var_map
529 .iter()
530 .map(|(var, &idx)| (idx, var.clone()))
531 .collect();
532
533 let mut solution_counts: HashMap<Vec<bool>, (f64, usize)> = HashMap::new();
534
535 if matrix_or_tensor.ndim() == 2 {
536 let q2 = matrix_or_tensor
538 .to_owned()
539 .into_dimensionality::<Ix2>()
540 .map_err(|e| SamplerError::InvalidParameter(format!("Array cast error: {e}")))?;
541 let n = q2.dim().0;
542 if n != q2.dim().1 {
543 return Err(SamplerError::InvalidParameter(
544 "QUBO matrix must be square".to_string(),
545 ));
546 }
547 let q_flat: Vec<f64> = q2
548 .as_slice()
549 .ok_or_else(|| {
550 SamplerError::InvalidParameter("Non-contiguous QUBO matrix".to_string())
551 })?
552 .to_vec();
553
554 for shot_idx in 0..shots {
555 let seed = match self.seed {
556 Some(s) => s.wrapping_add(shot_idx as u64),
557 None => {
558 let mut rng_tmp = thread_rng();
559 rng_tmp.random()
560 }
561 };
562 let mut rng = StdRng::seed_from_u64(seed);
563 let (best_state, best_energy) = self.run_single_qubo(&q_flat, n, &mut rng);
564
565 let entry = solution_counts
566 .entry(best_state)
567 .or_insert((best_energy, 0));
568 entry.1 += 1;
569 }
570 } else {
571 for shot_idx in 0..shots {
573 let seed = match self.seed {
574 Some(s) => s.wrapping_add(shot_idx as u64),
575 None => {
576 let mut rng_tmp = thread_rng();
577 rng_tmp.random()
578 }
579 };
580 let mut rng = StdRng::seed_from_u64(seed);
581 let (best_state, best_energy) =
582 self.run_single_hobo(matrix_or_tensor, n_vars, &mut rng);
583
584 let entry = solution_counts
585 .entry(best_state)
586 .or_insert((best_energy, 0));
587 entry.1 += 1;
588 }
589 }
590
591 let mut pairs: Vec<(Vec<bool>, SampleResult)> = solution_counts
593 .into_iter()
594 .map(|(state, (energy, count))| {
595 let assignments: HashMap<String, bool> = state
596 .iter()
597 .enumerate()
598 .filter_map(|(idx, &value)| {
599 idx_to_var.get(&idx).map(|name| (name.clone(), value))
600 })
601 .collect();
602 let result = SampleResult {
603 assignments,
604 energy,
605 occurrences: count,
606 };
607 (state, result)
608 })
609 .collect();
610
611 pairs.sort_by(|(state_a, a), (state_b, b)| {
613 a.energy
614 .partial_cmp(&b.energy)
615 .unwrap_or(std::cmp::Ordering::Equal)
616 .then_with(|| state_a.cmp(state_b))
617 });
618
619 let results: Vec<SampleResult> = pairs.into_iter().map(|(_, r)| r).collect();
620 Ok(results)
621 }
622}
623
624impl Sampler for TabuSampler {
625 fn run_qubo(
626 &self,
627 qubo: &(Array<f64, Ix2>, HashMap<String, usize>),
628 shots: usize,
629 ) -> SamplerResult<Vec<SampleResult>> {
630 self.run_generic(&qubo.0, &qubo.1, shots)
631 }
632
633 fn run_hobo(
634 &self,
635 hobo: &(ArrayD<f64>, HashMap<String, usize>),
636 shots: usize,
637 ) -> SamplerResult<Vec<SampleResult>> {
638 self.run_generic(&hobo.0, &hobo.1, shots)
639 }
640}
641
642#[cfg(test)]
643mod tests {
644 use super::*;
645 use scirs2_core::ndarray::Array2;
646
647 fn build_k3_maxcut_qubo() -> (Array2<f64>, HashMap<String, usize>) {
655 let mut q = Array2::<f64>::zeros((3, 3));
656 q[[0, 0]] = -2.0;
658 q[[1, 1]] = -2.0;
659 q[[2, 2]] = -2.0;
660 q[[0, 1]] = 2.0;
662 q[[0, 2]] = 2.0;
663 q[[1, 2]] = 2.0;
664 let mut var_map = HashMap::new();
667 var_map.insert("x0".to_string(), 0);
668 var_map.insert("x1".to_string(), 1);
669 var_map.insert("x2".to_string(), 2);
670
671 (q, var_map)
672 }
673
674 #[test]
675 fn test_tabu_3var_maxcut() {
676 let (q, var_map) = build_k3_maxcut_qubo();
677 let sampler = TabuSampler::new()
678 .with_seed(42)
679 .with_max_iter(200)
680 .with_tenure(5);
681
682 let results = sampler
683 .run_qubo(&(q, var_map), 50)
684 .expect("Tabu run_qubo failed");
685
686 assert!(!results.is_empty(), "Expected non-empty results");
687 let best_energy = results[0].energy;
689 assert!(
690 best_energy <= -2.0 + 1e-9,
691 "Expected optimal energy <= -2.0, got {best_energy}"
692 );
693 }
694
695 #[test]
696 fn test_tabu_determinism() {
697 let (q, var_map) = build_k3_maxcut_qubo();
698
699 let s1 = TabuSampler::new().with_seed(42);
700 let s2 = TabuSampler::new().with_seed(42);
701
702 let r1 = s1
703 .run_qubo(&(q.clone(), var_map.clone()), 10)
704 .expect("Run 1 failed");
705 let r2 = s2.run_qubo(&(q, var_map), 10).expect("Run 2 failed");
706
707 assert_eq!(r1.len(), r2.len(), "Result lengths differ");
708 for (a, b) in r1.iter().zip(r2.iter()) {
709 assert!(
710 (a.energy - b.energy).abs() < 1e-12,
711 "Energies differ: {} vs {}",
712 a.energy,
713 b.energy
714 );
715 assert_eq!(
716 a.assignments, b.assignments,
717 "Assignments differ for same seed"
718 );
719 }
720 }
721
722 #[test]
723 fn test_tabu_hobo_smoke() {
724 let mut q = Array2::<f64>::zeros((2, 2));
726 q[[0, 0]] = -1.0;
727 q[[1, 1]] = -1.0;
728
729 let mut var_map = HashMap::new();
730 var_map.insert("a".to_string(), 0);
731 var_map.insert("b".to_string(), 1);
732
733 let sampler = TabuSampler::new().with_seed(7);
734 let q_dyn = q.into_dyn();
735 let results = sampler
736 .run_hobo(&(q_dyn, var_map), 10)
737 .expect("HOBO run failed");
738 assert!(!results.is_empty());
739 assert!(results[0].energy <= -2.0 + 1e-9);
740 }
741
742 #[test]
743 fn test_tabu_aspiration() {
744 let (q, var_map) = build_k3_maxcut_qubo();
745 let with_aspiration = TabuSampler::new()
747 .with_seed(100)
748 .with_aspiration(true)
749 .with_tenure(100) .with_max_iter(500);
751 let without_aspiration = TabuSampler::new()
752 .with_seed(100)
753 .with_aspiration(false)
754 .with_max_iter(500);
755
756 let r1 = with_aspiration
757 .run_qubo(&(q.clone(), var_map.clone()), 20)
758 .expect("With aspiration failed");
759 let r2 = without_aspiration
760 .run_qubo(&(q, var_map), 20)
761 .expect("Without aspiration failed");
762
763 assert!(!r1.is_empty());
764 assert!(!r2.is_empty());
765 assert!(r1[0].energy <= -2.0 + 1e-9 || r2[0].energy <= -2.0 + 1e-9);
767 }
768}