1use crate::error::{PgmError, Result};
42use crate::graph::FactorGraph;
43use quantrs2_sim::circuit_optimizer::{Circuit, Gate, GateType};
44use scirs2_core::ndarray::{Array1, Array2};
45use serde::{Deserialize, Serialize};
46use std::collections::HashMap;
47use std::f64::consts::PI;
48use tensorlogic_ir::TLExpr;
49
50#[derive(Debug, Clone, Serialize, Deserialize)]
57pub struct QUBOProblem {
58 pub num_variables: usize,
60 pub variable_names: Vec<String>,
62 pub quadratic: Array2<f64>,
64 pub linear: Array1<f64>,
66 pub offset: f64,
68}
69
70impl QUBOProblem {
71 pub fn new(num_variables: usize) -> Self {
73 Self {
74 num_variables,
75 variable_names: (0..num_variables).map(|i| format!("x_{}", i)).collect(),
76 quadratic: Array2::zeros((num_variables, num_variables)),
77 linear: Array1::zeros(num_variables),
78 offset: 0.0,
79 }
80 }
81
82 pub fn with_names(variable_names: Vec<String>) -> Self {
84 let n = variable_names.len();
85 Self {
86 num_variables: n,
87 variable_names,
88 quadratic: Array2::zeros((n, n)),
89 linear: Array1::zeros(n),
90 offset: 0.0,
91 }
92 }
93
94 pub fn set_linear(&mut self, i: usize, value: f64) {
96 if i < self.num_variables {
97 self.linear[i] = value;
98 }
99 }
100
101 pub fn set_quadratic(&mut self, i: usize, j: usize, value: f64) {
103 if i < self.num_variables && j < self.num_variables {
104 if i <= j {
105 self.quadratic[[i, j]] = value;
106 } else {
107 self.quadratic[[j, i]] = value;
108 }
109 }
110 }
111
112 pub fn add_linear(&mut self, i: usize, value: f64) {
114 if i < self.num_variables {
115 self.linear[i] += value;
116 }
117 }
118
119 pub fn add_quadratic(&mut self, i: usize, j: usize, value: f64) {
121 if i < self.num_variables && j < self.num_variables {
122 if i <= j {
123 self.quadratic[[i, j]] += value;
124 } else {
125 self.quadratic[[j, i]] += value;
126 }
127 }
128 }
129
130 pub fn evaluate(&self, assignment: &[usize]) -> f64 {
132 let mut value = self.offset;
133
134 for (i, &xi) in assignment.iter().enumerate() {
136 if i < self.num_variables {
137 value += self.linear[i] * xi as f64;
138 }
139 }
140
141 for i in 0..self.num_variables {
143 for j in i..self.num_variables {
144 let xi = if i < assignment.len() {
145 assignment[i]
146 } else {
147 0
148 };
149 let xj = if j < assignment.len() {
150 assignment[j]
151 } else {
152 0
153 };
154 value += self.quadratic[[i, j]] * (xi * xj) as f64;
155 }
156 }
157
158 value
159 }
160
161 pub fn to_ising(&self) -> IsingModel {
168 let n = self.num_variables;
169 let mut h = Array1::zeros(n);
170 let mut j_matrix = Array2::zeros((n, n));
171 let mut offset = self.offset;
172
173 for i in 0..n {
175 h[i] = self.linear[i] / 2.0;
176 offset += self.linear[i] / 2.0;
177 }
178
179 for i in 0..n {
181 for j in i..n {
182 let qij = self.quadratic[[i, j]];
183 if i == j {
184 h[i] += qij / 2.0;
186 offset += qij / 2.0;
187 } else {
188 j_matrix[[i, j]] = qij / 4.0;
190 h[i] += qij / 4.0;
191 h[j] += qij / 4.0;
192 offset += qij / 4.0;
193 }
194 }
195 }
196
197 IsingModel {
198 num_spins: n,
199 h,
200 j: j_matrix,
201 offset,
202 }
203 }
204
205 pub fn variable_index(&self, name: &str) -> Option<usize> {
207 self.variable_names.iter().position(|n| n == name)
208 }
209}
210
211#[derive(Debug, Clone)]
215pub struct IsingModel {
216 pub num_spins: usize,
218 pub h: Array1<f64>,
220 pub j: Array2<f64>,
222 pub offset: f64,
224}
225
226impl IsingModel {
227 pub fn evaluate(&self, spins: &[i32]) -> f64 {
231 let mut energy = self.offset;
232
233 for i in 0..self.num_spins {
235 if i < spins.len() {
236 energy += self.h[i] * spins[i] as f64;
237 }
238 }
239
240 for i in 0..self.num_spins {
242 for j in (i + 1)..self.num_spins {
243 if i < spins.len() && j < spins.len() {
244 energy += self.j[[i, j]] * (spins[i] * spins[j]) as f64;
245 }
246 }
247 }
248
249 energy
250 }
251}
252
253pub fn constraint_to_qubo(constraints: &[TLExpr], num_variables: usize) -> Result<QUBOProblem> {
272 let mut qubo = QUBOProblem::new(num_variables);
273 let mut var_map: HashMap<String, usize> = HashMap::new();
274 let mut next_idx = 0;
275
276 for expr in constraints {
278 collect_variables(expr, &mut var_map, &mut next_idx, num_variables)?;
279 }
280
281 qubo.variable_names = vec![String::new(); num_variables];
283 for (name, &idx) in &var_map {
284 if idx < num_variables {
285 qubo.variable_names[idx] = name.clone();
286 }
287 }
288
289 for expr in constraints {
291 add_constraint_penalty(expr, &var_map, &mut qubo)?;
292 }
293
294 Ok(qubo)
295}
296
297fn collect_variables(
299 expr: &TLExpr,
300 var_map: &mut HashMap<String, usize>,
301 next_idx: &mut usize,
302 max_vars: usize,
303) -> Result<()> {
304 match expr {
305 TLExpr::Pred { args, .. } => {
306 for term in args {
307 if let tensorlogic_ir::Term::Var(v) = term {
308 if !var_map.contains_key(v) && *next_idx < max_vars {
309 var_map.insert(v.clone(), *next_idx);
310 *next_idx += 1;
311 }
312 }
313 }
314 }
315 TLExpr::And(left, right) | TLExpr::Or(left, right) | TLExpr::Imply(left, right) => {
316 collect_variables(left, var_map, next_idx, max_vars)?;
317 collect_variables(right, var_map, next_idx, max_vars)?;
318 }
319 TLExpr::Not(inner) => {
320 collect_variables(inner, var_map, next_idx, max_vars)?;
321 }
322 TLExpr::Exists { body, .. } | TLExpr::ForAll { body, .. } => {
323 collect_variables(body, var_map, next_idx, max_vars)?;
324 }
325 _ => {}
326 }
327 Ok(())
328}
329
330fn add_constraint_penalty(
332 expr: &TLExpr,
333 var_map: &HashMap<String, usize>,
334 qubo: &mut QUBOProblem,
335) -> Result<()> {
336 match expr {
337 TLExpr::Pred { name, args } => {
338 let mut var_indices = Vec::new();
340 for term in args {
341 if let tensorlogic_ir::Term::Var(v) = term {
342 if let Some(&idx) = var_map.get(v) {
343 var_indices.push(idx);
344 }
345 }
346 }
347
348 match name.as_str() {
350 "edge" | "connected" | "equal" if var_indices.len() >= 2 => {
351 let i = var_indices[0];
353 let j = var_indices[1];
354 qubo.add_linear(i, 1.0);
355 qubo.add_linear(j, 1.0);
356 qubo.add_quadratic(i, j, -2.0);
357 }
358 "conflict" | "different" | "not_equal" if var_indices.len() >= 2 => {
359 let i = var_indices[0];
361 let j = var_indices[1];
362 qubo.add_quadratic(i, j, 1.0);
363 }
364 "select" | "chosen" if !var_indices.is_empty() => {
365 let i = var_indices[0];
367 qubo.add_linear(i, -1.0);
368 }
369 "exclude" | "reject" if !var_indices.is_empty() => {
370 let i = var_indices[0];
372 qubo.add_linear(i, 1.0);
373 }
374 _ => {
375 for i in 0..var_indices.len() {
377 for j in (i + 1)..var_indices.len() {
378 qubo.add_quadratic(var_indices[i], var_indices[j], 0.5);
379 }
380 }
381 }
382 }
383 }
384 TLExpr::And(left, right) => {
385 add_constraint_penalty(left, var_map, qubo)?;
386 add_constraint_penalty(right, var_map, qubo)?;
387 }
388 TLExpr::Not(inner) => {
389 add_constraint_penalty(inner, var_map, qubo)?;
392 }
393 _ => {
394 }
396 }
397 Ok(())
398}
399
400pub fn tlexpr_to_qaoa_circuit(
417 expr: &TLExpr,
418 num_layers: usize,
419 num_qubits: usize,
420) -> Result<Circuit> {
421 let constraints = vec![expr.clone()];
423 let qubo = constraint_to_qubo(&constraints, num_qubits)?;
424
425 let builder = QuantumCircuitBuilder::new(num_qubits);
427 builder.build_qaoa_circuit(&qubo, num_layers)
428}
429
430pub fn factor_graph_to_qubo(graph: &FactorGraph) -> Result<QUBOProblem> {
434 let num_vars = graph.num_variables();
435 let mut qubo = QUBOProblem::new(num_vars);
436
437 let mut var_to_idx: HashMap<String, usize> = HashMap::new();
439 for (idx, var_name) in graph.variable_names().enumerate() {
440 var_to_idx.insert(var_name.clone(), idx);
441 if idx < qubo.variable_names.len() {
442 qubo.variable_names[idx] = var_name.clone();
443 }
444 }
445
446 for factor in graph.factors() {
448 let factor_vars: Vec<usize> = factor
449 .variables
450 .iter()
451 .filter_map(|v| var_to_idx.get(v).copied())
452 .collect();
453
454 if factor_vars.len() == 1 {
456 let i = factor_vars[0];
458 if factor.values.len() >= 2 {
459 let penalty = -(factor.values[[1]].ln() - factor.values[[0]].ln());
461 if penalty.is_finite() {
462 qubo.add_linear(i, penalty);
463 }
464 }
465 } else if factor_vars.len() == 2 {
466 let i = factor_vars[0];
468 let j = factor_vars[1];
469 if factor.values.len() >= 4 {
470 let f00 = factor.values[[0, 0]].max(1e-10);
473 let f01 = factor.values[[0, 1]].max(1e-10);
474 let f10 = factor.values[[1, 0]].max(1e-10);
475 let f11 = factor.values[[1, 1]].max(1e-10);
476
477 let coupling = ((f00 * f11).ln() - (f01 * f10).ln()) / 4.0;
478 if coupling.is_finite() {
479 qubo.add_quadratic(i, j, -coupling);
480 }
481 }
482 }
483 }
485
486 Ok(qubo)
487}
488
489pub struct QuantumCircuitBuilder {
494 num_qubits: usize,
496 default_params: Vec<f64>,
498}
499
500impl QuantumCircuitBuilder {
501 pub fn new(num_qubits: usize) -> Self {
503 Self {
504 num_qubits,
505 default_params: vec![PI / 4.0, PI / 4.0], }
507 }
508
509 pub fn with_params(mut self, params: Vec<f64>) -> Self {
511 self.default_params = params;
512 self
513 }
514
515 pub fn build_initial_state(&self) -> Result<Circuit> {
517 let mut circuit = Circuit::new(self.num_qubits);
518
519 for qubit in 0..self.num_qubits {
521 let _ = circuit.add_gate(Gate::new(GateType::H, vec![qubit]));
522 }
523
524 Ok(circuit)
525 }
526
527 pub fn build_qaoa_circuit(&self, qubo: &QUBOProblem, num_layers: usize) -> Result<Circuit> {
542 if qubo.num_variables > self.num_qubits {
543 return Err(PgmError::InvalidGraph(format!(
544 "QUBO has {} variables but circuit has {} qubits",
545 qubo.num_variables, self.num_qubits
546 )));
547 }
548
549 let ising = qubo.to_ising();
551
552 let mut circuit = Circuit::new(self.num_qubits);
553
554 for qubit in 0..self.num_qubits {
556 let _ = circuit.add_gate(Gate::new(GateType::H, vec![qubit]));
557 }
558
559 for layer in 0..num_layers {
561 let gamma = if layer * 2 < self.default_params.len() {
563 self.default_params[layer * 2]
564 } else {
565 PI / (4.0 * (layer + 1) as f64)
566 };
567
568 let beta = if layer * 2 + 1 < self.default_params.len() {
569 self.default_params[layer * 2 + 1]
570 } else {
571 PI / (4.0 * (layer + 1) as f64)
572 };
573
574 self.add_problem_unitary(&mut circuit, &ising, gamma);
576
577 self.add_mixer_unitary(&mut circuit, beta);
579 }
580
581 Ok(circuit)
582 }
583
584 fn add_problem_unitary(&self, circuit: &mut Circuit, ising: &IsingModel, gamma: f64) {
586 let n = ising.num_spins.min(self.num_qubits);
587
588 for i in 0..n {
590 if ising.h[i].abs() > 1e-10 {
591 let angle = 2.0 * gamma * ising.h[i];
592 let _ = circuit.add_gate(Gate::with_parameters(GateType::RZ, vec![i], vec![angle]));
593 }
594 }
595
596 for i in 0..n {
598 for j in (i + 1)..n {
599 let jij = ising.j[[i, j]];
600 if jij.abs() > 1e-10 {
601 self.add_rzz_gate(circuit, i, j, 2.0 * gamma * jij);
602 }
603 }
604 }
605 }
606
607 fn add_mixer_unitary(&self, circuit: &mut Circuit, beta: f64) {
609 for qubit in 0..self.num_qubits {
611 let _ = circuit.add_gate(Gate::with_parameters(
612 GateType::RX,
613 vec![qubit],
614 vec![2.0 * beta],
615 ));
616 }
617 }
618
619 fn add_rzz_gate(&self, circuit: &mut Circuit, qubit1: usize, qubit2: usize, angle: f64) {
624 let _ = circuit.add_gate(Gate::new(GateType::CNOT, vec![qubit1, qubit2]));
625 let _ = circuit.add_gate(Gate::with_parameters(
626 GateType::RZ,
627 vec![qubit2],
628 vec![angle],
629 ));
630 let _ = circuit.add_gate(Gate::new(GateType::CNOT, vec![qubit1, qubit2]));
631 }
632
633 pub fn build_hardware_efficient_ansatz(&self, num_layers: usize) -> Result<Circuit> {
638 let mut circuit = Circuit::new(self.num_qubits);
639
640 for _layer in 0..num_layers {
641 for qubit in 0..self.num_qubits {
643 let _ = circuit.add_gate(Gate::with_parameters(
644 GateType::RY,
645 vec![qubit],
646 vec![PI / 4.0],
647 ));
648 let _ = circuit.add_gate(Gate::with_parameters(
649 GateType::RZ,
650 vec![qubit],
651 vec![PI / 4.0],
652 ));
653 }
654
655 for qubit in 0..(self.num_qubits - 1) {
657 let _ = circuit.add_gate(Gate::new(GateType::CNOT, vec![qubit, qubit + 1]));
658 }
659 }
660
661 for qubit in 0..self.num_qubits {
663 let _ = circuit.add_gate(Gate::with_parameters(
664 GateType::RY,
665 vec![qubit],
666 vec![PI / 4.0],
667 ));
668 }
669
670 Ok(circuit)
671 }
672}
673
674#[derive(Debug, Clone)]
676pub struct QAOAResult {
677 pub gamma: Vec<f64>,
679 pub beta: Vec<f64>,
681 pub best_solution: Vec<usize>,
683 pub best_value: f64,
685 pub iterations: usize,
687}
688
689#[derive(Debug, Clone)]
691pub struct QAOAConfig {
692 pub num_layers: usize,
694 pub num_shots: usize,
696 pub max_iterations: usize,
698 pub tolerance: f64,
700}
701
702impl Default for QAOAConfig {
703 fn default() -> Self {
704 Self {
705 num_layers: 2,
706 num_shots: 1000,
707 max_iterations: 100,
708 tolerance: 1e-6,
709 }
710 }
711}
712
713impl QAOAConfig {
714 pub fn new(num_layers: usize) -> Self {
716 Self {
717 num_layers,
718 ..Default::default()
719 }
720 }
721
722 pub fn with_shots(mut self, shots: usize) -> Self {
724 self.num_shots = shots;
725 self
726 }
727
728 pub fn with_max_iterations(mut self, iterations: usize) -> Self {
730 self.max_iterations = iterations;
731 self
732 }
733}
734
735#[cfg(test)]
736mod tests {
737 use super::*;
738 use approx::assert_abs_diff_eq;
739
740 #[test]
741 fn test_qubo_creation() {
742 let qubo = QUBOProblem::new(3);
743 assert_eq!(qubo.num_variables, 3);
744 assert_eq!(qubo.variable_names.len(), 3);
745 }
746
747 #[test]
748 fn test_qubo_with_names() {
749 let qubo = QUBOProblem::with_names(vec!["x".to_string(), "y".to_string(), "z".to_string()]);
750 assert_eq!(qubo.num_variables, 3);
751 assert_eq!(qubo.variable_names[0], "x");
752 assert_eq!(qubo.variable_index("y"), Some(1));
753 }
754
755 #[test]
756 fn test_qubo_evaluation() {
757 let mut qubo = QUBOProblem::new(2);
758 qubo.set_linear(0, 1.0);
759 qubo.set_linear(1, 2.0);
760 qubo.set_quadratic(0, 1, 3.0);
761
762 assert_abs_diff_eq!(qubo.evaluate(&[0, 0]), 0.0);
764
765 assert_abs_diff_eq!(qubo.evaluate(&[1, 0]), 1.0);
767
768 assert_abs_diff_eq!(qubo.evaluate(&[0, 1]), 2.0);
770
771 assert_abs_diff_eq!(qubo.evaluate(&[1, 1]), 6.0);
773 }
774
775 #[test]
776 fn test_qubo_to_ising() {
777 let mut qubo = QUBOProblem::new(2);
778 qubo.set_linear(0, 1.0);
779 qubo.set_quadratic(0, 1, 2.0);
780
781 let ising = qubo.to_ising();
782 assert_eq!(ising.num_spins, 2);
783
784 let qubo_00 = qubo.evaluate(&[0, 0]);
787 let qubo_11 = qubo.evaluate(&[1, 1]);
788
789 let ising_mm = ising.evaluate(&[-1, -1]);
790 let ising_pp = ising.evaluate(&[1, 1]);
791
792 assert_abs_diff_eq!(qubo_11 - qubo_00, ising_pp - ising_mm, epsilon = 1e-6);
794 }
795
796 #[test]
797 fn test_constraint_to_qubo() {
798 use tensorlogic_ir::Term;
799
800 let constraints = vec![TLExpr::pred("edge", vec![Term::var("x"), Term::var("y")])];
801
802 let qubo = constraint_to_qubo(&constraints, 4).ok();
803 assert!(qubo.is_some());
804
805 let qubo = qubo.expect("QUBO creation failed");
806 assert_eq!(qubo.num_variables, 4);
807 }
808
809 #[test]
810 fn test_circuit_builder_initial_state() {
811 let builder = QuantumCircuitBuilder::new(3);
812 let circuit = builder.build_initial_state();
813
814 assert!(circuit.is_ok());
815 let circuit = circuit.expect("Circuit creation failed");
816 assert_eq!(circuit.n_qubits, 3);
817 }
818
819 #[test]
820 fn test_qaoa_circuit_creation() {
821 let mut qubo = QUBOProblem::new(2);
822 qubo.set_quadratic(0, 1, 1.0);
823
824 let builder = QuantumCircuitBuilder::new(2);
825 let circuit = builder.build_qaoa_circuit(&qubo, 1);
826
827 assert!(circuit.is_ok());
828 let circuit = circuit.expect("Circuit creation failed");
829 assert_eq!(circuit.n_qubits, 2);
830 }
831
832 #[test]
833 fn test_hardware_efficient_ansatz() {
834 let builder = QuantumCircuitBuilder::new(3);
835 let circuit = builder.build_hardware_efficient_ansatz(2);
836
837 assert!(circuit.is_ok());
838 let circuit = circuit.expect("Circuit creation failed");
839 assert_eq!(circuit.n_qubits, 3);
840 }
841
842 #[test]
843 fn test_qaoa_config() {
844 let config = QAOAConfig::new(3).with_shots(2000).with_max_iterations(50);
845
846 assert_eq!(config.num_layers, 3);
847 assert_eq!(config.num_shots, 2000);
848 assert_eq!(config.max_iterations, 50);
849 }
850
851 #[test]
852 fn test_tlexpr_to_qaoa() {
853 use tensorlogic_ir::Term;
854
855 let expr = TLExpr::pred("conflict", vec![Term::var("a"), Term::var("b")]);
856
857 let circuit = tlexpr_to_qaoa_circuit(&expr, 1, 4);
858 assert!(circuit.is_ok());
859 }
860}