1use scirs2_core::ndarray::{Array1, Array2};
7use std::collections::HashMap;
8use std::f64::consts::PI;
9
10use crate::error::{MLError, Result};
11use quantrs2_circuit::prelude::*;
12use quantrs2_core::gate::GateOp;
13
14#[derive(Debug, Clone)]
16pub struct DifferentiableParam {
17 pub name: String,
19 pub value: f64,
21 pub gradient: f64,
23 pub requires_grad: bool,
25}
26
27impl DifferentiableParam {
28 pub fn new(name: impl Into<String>, value: f64) -> Self {
30 Self {
31 name: name.into(),
32 value,
33 gradient: 0.0,
34 requires_grad: true,
35 }
36 }
37
38 pub fn constant(name: impl Into<String>, value: f64) -> Self {
40 Self {
41 name: name.into(),
42 value,
43 gradient: 0.0,
44 requires_grad: false,
45 }
46 }
47}
48
49#[derive(Debug, Clone)]
51pub enum ComputationNode {
52 Parameter(String),
54 Constant(f64),
56 Add(Box<ComputationNode>, Box<ComputationNode>),
58 Mul(Box<ComputationNode>, Box<ComputationNode>),
60 Sin(Box<ComputationNode>),
62 Cos(Box<ComputationNode>),
64 Exp(Box<ComputationNode>),
66 Expectation {
68 circuit_params: Vec<String>,
69 observable: String,
70 },
71}
72
73pub struct AutoDiff {
75 parameters: HashMap<String, DifferentiableParam>,
77 graph: Option<ComputationNode>,
79 forward_cache: HashMap<String, f64>,
81 executor: Option<Box<dyn Fn(&[f64], &str) -> f64>>,
89}
90
91impl AutoDiff {
92 pub fn new() -> Self {
94 Self {
95 parameters: HashMap::new(),
96 graph: None,
97 forward_cache: HashMap::new(),
98 executor: None,
99 }
100 }
101
102 pub fn with_executor<F>(mut self, executor: F) -> Self
107 where
108 F: Fn(&[f64], &str) -> f64 + 'static,
109 {
110 self.executor = Some(Box::new(executor));
111 self
112 }
113
114 pub fn register_parameter(&mut self, param: DifferentiableParam) {
116 self.parameters.insert(param.name.clone(), param);
117 }
118
119 pub fn set_graph(&mut self, graph: ComputationNode) {
121 self.graph = Some(graph);
122 }
123
124 pub fn forward(&mut self) -> Result<f64> {
126 self.forward_cache.clear();
127
128 if let Some(graph) = self.graph.clone() {
129 self.evaluate_node(&graph)
130 } else {
131 Err(MLError::InvalidConfiguration(
132 "No computation graph set".to_string(),
133 ))
134 }
135 }
136
137 pub fn backward(&mut self, loss_gradient: f64) -> Result<()> {
139 for param in self.parameters.values_mut() {
141 param.gradient = 0.0;
142 }
143
144 if let Some(graph) = self.graph.clone() {
145 self.backpropagate(&graph, loss_gradient)?;
146 }
147
148 Ok(())
149 }
150
151 fn evaluate_node(&mut self, node: &ComputationNode) -> Result<f64> {
153 match node {
154 ComputationNode::Parameter(name) => {
155 self.parameters.get(name).map(|p| p.value).ok_or_else(|| {
156 MLError::InvalidConfiguration(format!("Unknown parameter: {}", name))
157 })
158 }
159 ComputationNode::Constant(value) => Ok(*value),
160 ComputationNode::Add(left, right) => {
161 let l = self.evaluate_node(left)?;
162 let r = self.evaluate_node(right)?;
163 Ok(l + r)
164 }
165 ComputationNode::Mul(left, right) => {
166 let l = self.evaluate_node(left)?;
167 let r = self.evaluate_node(right)?;
168 Ok(l * r)
169 }
170 ComputationNode::Sin(inner) => {
171 let x = self.evaluate_node(inner)?;
172 Ok(x.sin())
173 }
174 ComputationNode::Cos(inner) => {
175 let x = self.evaluate_node(inner)?;
176 Ok(x.cos())
177 }
178 ComputationNode::Exp(inner) => {
179 let x = self.evaluate_node(inner)?;
180 Ok(x.exp())
181 }
182 ComputationNode::Expectation {
183 circuit_params,
184 observable,
185 } => {
186 let values = self.circuit_param_values(circuit_params)?;
187 let executor = self.executor.as_ref().ok_or_else(|| {
188 MLError::NotSupported(
189 "ComputationNode::Expectation requires a circuit executor; call \
190 AutoDiff::with_executor() before forward()/backward()"
191 .to_string(),
192 )
193 })?;
194 Ok(executor(&values, observable))
195 }
196 }
197 }
198
199 fn circuit_param_values(&self, circuit_params: &[String]) -> Result<Vec<f64>> {
202 circuit_params
203 .iter()
204 .map(|name| {
205 self.parameters.get(name).map(|p| p.value).ok_or_else(|| {
206 MLError::InvalidConfiguration(format!("Unknown parameter: {}", name))
207 })
208 })
209 .collect()
210 }
211
212 fn backpropagate(&mut self, node: &ComputationNode, grad: f64) -> Result<()> {
214 match node {
215 ComputationNode::Parameter(name) => {
216 if let Some(param) = self.parameters.get_mut(name) {
217 if param.requires_grad {
218 param.gradient += grad;
219 }
220 }
221 }
222 ComputationNode::Constant(_) => {
223 }
225 ComputationNode::Add(left, right) => {
226 self.backpropagate(left, grad)?;
228 self.backpropagate(right, grad)?;
229 }
230 ComputationNode::Mul(left, right) => {
231 let l_val = self.evaluate_node(left)?;
233 let r_val = self.evaluate_node(right)?;
234 self.backpropagate(left, grad * r_val)?;
235 self.backpropagate(right, grad * l_val)?;
236 }
237 ComputationNode::Sin(inner) => {
238 let x = self.evaluate_node(inner)?;
240 self.backpropagate(inner, grad * x.cos())?;
241 }
242 ComputationNode::Cos(inner) => {
243 let x = self.evaluate_node(inner)?;
245 self.backpropagate(inner, grad * (-x.sin()))?;
246 }
247 ComputationNode::Exp(inner) => {
248 let x = self.evaluate_node(inner)?;
250 self.backpropagate(inner, grad * x.exp())?;
251 }
252 ComputationNode::Expectation {
253 circuit_params,
254 observable,
255 } => {
256 for (index, param_name) in circuit_params.iter().enumerate() {
259 let shift_grad =
260 self.parameter_shift_gradient(circuit_params, observable, index, PI / 2.0)?;
261 if let Some(param) = self.parameters.get_mut(param_name) {
262 if param.requires_grad {
263 param.gradient += grad * shift_grad;
264 }
265 }
266 }
267 }
268 }
269 Ok(())
270 }
271
272 fn parameter_shift_gradient(
277 &self,
278 circuit_params: &[String],
279 observable: &str,
280 index: usize,
281 shift: f64,
282 ) -> Result<f64> {
283 let executor = self.executor.as_ref().ok_or_else(|| {
284 MLError::NotSupported(
285 "parameter_shift_gradient requires a circuit executor; call \
286 AutoDiff::with_executor() before forward()/backward()"
287 .to_string(),
288 )
289 })?;
290 let mut values = self.circuit_param_values(circuit_params)?;
291 if index >= values.len() {
292 return Err(MLError::InvalidParameter(format!(
293 "parameter index {index} out of range for {} circuit parameters",
294 values.len()
295 )));
296 }
297
298 let original = values[index];
299 values[index] = original + shift;
300 let plus = executor(&values, observable);
301 values[index] = original - shift;
302 let minus = executor(&values, observable);
303
304 Ok((plus - minus) / (2.0 * shift.sin()))
305 }
306
307 pub fn gradients(&self) -> HashMap<String, f64> {
309 self.parameters
310 .iter()
311 .filter(|(_, p)| p.requires_grad)
312 .map(|(name, param)| (name.clone(), param.gradient))
313 .collect()
314 }
315
316 pub fn update_parameters(&mut self, learning_rate: f64) {
318 for param in self.parameters.values_mut() {
319 if param.requires_grad {
320 param.value -= learning_rate * param.gradient;
321 }
322 }
323 }
324}
325
326pub struct QuantumAutoDiff {
328 autodiff: AutoDiff,
330 executor: Box<dyn Fn(&[f64]) -> f64>,
332}
333
334impl QuantumAutoDiff {
335 pub fn new<F>(executor: F) -> Self
337 where
338 F: Fn(&[f64]) -> f64 + 'static,
339 {
340 Self {
341 autodiff: AutoDiff::new(),
342 executor: Box::new(executor),
343 }
344 }
345
346 pub fn parameter_shift_gradients(&self, params: &[f64], shift: f64) -> Result<Vec<f64>> {
348 let mut gradients = vec![0.0; params.len()];
349
350 for (i, _) in params.iter().enumerate() {
351 let mut params_plus = params.to_vec();
353 params_plus[i] += shift;
354 let val_plus = (self.executor)(¶ms_plus);
355
356 let mut params_minus = params.to_vec();
358 params_minus[i] -= shift;
359 let val_minus = (self.executor)(¶ms_minus);
360
361 gradients[i] = (val_plus - val_minus) / (2.0 * shift.sin());
363 }
364
365 Ok(gradients)
366 }
367
368 pub fn natural_gradients(
370 &self,
371 params: &[f64],
372 gradients: &[f64],
373 regularization: f64,
374 ) -> Result<Vec<f64>> {
375 let n = params.len();
376 let mut fisher = Array2::<f64>::zeros((n, n));
377
378 for i in 0..n {
380 for j in 0..n {
381 fisher[[i, j]] = self.compute_fisher_element(params, i, j)?;
382 }
383 }
384
385 for i in 0..n {
387 fisher[[i, i]] += regularization;
388 }
389
390 self.solve_linear_system(&fisher, gradients)
392 }
393
394 fn compute_fisher_element(&self, params: &[f64], i: usize, j: usize) -> Result<f64> {
399 let shift = PI / 2.0;
400
401 let mut p_pp = params.to_vec();
402 let mut p_pm = params.to_vec();
403 let mut p_mp = params.to_vec();
404 let mut p_mm = params.to_vec();
405
406 p_pp[i] += shift;
407 p_pp[j] += shift;
408
409 p_pm[i] += shift;
410 p_pm[j] -= shift;
411
412 p_mp[i] -= shift;
413 p_mp[j] += shift;
414
415 p_mm[i] -= shift;
416 p_mm[j] -= shift;
417
418 let e_pp = (self.executor)(&p_pp);
419 let e_pm = (self.executor)(&p_pm);
420 let e_mp = (self.executor)(&p_mp);
421 let e_mm = (self.executor)(&p_mm);
422
423 Ok((e_pp - e_pm - e_mp + e_mm) / 4.0)
424 }
425
426 fn solve_linear_system(&self, matrix: &Array2<f64>, rhs: &[f64]) -> Result<Vec<f64>> {
431 let n = rhs.len();
432 if matrix.nrows() != n || matrix.ncols() != n {
433 return Err(MLError::DimensionMismatch(format!(
434 "Matrix ({} x {}) incompatible with rhs length {}",
435 matrix.nrows(),
436 matrix.ncols(),
437 n
438 )));
439 }
440
441 let mut a: Vec<Vec<f64>> = (0..n)
443 .map(|i| {
444 let mut row: Vec<f64> = (0..n).map(|j| matrix[[i, j]]).collect();
445 row.push(rhs[i]);
446 row
447 })
448 .collect();
449
450 for k in 0..n {
452 let mut max_val = a[k][k].abs();
454 let mut max_idx = k;
455 for row in (k + 1)..n {
456 let val = a[row][k].abs();
457 if val > max_val {
458 max_val = val;
459 max_idx = row;
460 }
461 }
462
463 if max_val < 1e-12 {
464 return Err(MLError::NumericalError(format!(
465 "Singular matrix: |pivot| = {:.2e} < 1e-12 at column {}",
466 max_val, k
467 )));
468 }
469
470 if max_idx != k {
472 a.swap(k, max_idx);
473 }
474
475 let pivot = a[k][k];
476
477 for i in (k + 1)..n {
479 let factor = a[i][k] / pivot;
480 for col in k..=n {
481 let sub = factor * a[k][col];
482 a[i][col] -= sub;
483 }
484 }
485 }
486
487 let mut x = vec![0.0_f64; n];
489 for i in (0..n).rev() {
490 let mut sum = a[i][n]; for j in (i + 1)..n {
492 sum -= a[i][j] * x[j];
493 }
494 x[i] = sum / a[i][i];
495 }
496
497 Ok(x)
498 }
499}
500
501#[derive(Debug, Clone)]
503pub struct GradientTape {
504 operations: Vec<Operation>,
506 variables: HashMap<String, f64>,
508}
509
510#[derive(Debug, Clone)]
512enum Operation {
513 Assign { var: String, value: f64 },
515 Add {
517 result: String,
518 left: String,
519 right: String,
520 },
521 Mul {
523 result: String,
524 left: String,
525 right: String,
526 },
527 Quantum { result: String, params: Vec<String> },
529}
530
531impl GradientTape {
532 pub fn new() -> Self {
534 Self {
535 operations: Vec::new(),
536 variables: HashMap::new(),
537 }
538 }
539
540 pub fn variable(&mut self, name: impl Into<String>, value: f64) -> String {
542 let name = name.into();
543 self.variables.insert(name.clone(), value);
544 self.operations.push(Operation::Assign {
545 var: name.clone(),
546 value,
547 });
548 name
549 }
550
551 pub fn add(&mut self, left: &str, right: &str) -> String {
553 let result = format!("tmp_{}", self.operations.len());
554 let left_val = self.variables[left];
555 let right_val = self.variables[right];
556 self.variables.insert(result.clone(), left_val + right_val);
557 self.operations.push(Operation::Add {
558 result: result.clone(),
559 left: left.to_string(),
560 right: right.to_string(),
561 });
562 result
563 }
564
565 pub fn mul(&mut self, left: &str, right: &str) -> String {
567 let result = format!("tmp_{}", self.operations.len());
568 let left_val = self.variables[left];
569 let right_val = self.variables[right];
570 self.variables.insert(result.clone(), left_val * right_val);
571 self.operations.push(Operation::Mul {
572 result: result.clone(),
573 left: left.to_string(),
574 right: right.to_string(),
575 });
576 result
577 }
578
579 pub fn gradient(&self, output: &str, inputs: &[&str]) -> HashMap<String, f64> {
581 let mut gradients: HashMap<String, f64> = HashMap::new();
582
583 gradients.insert(output.to_string(), 1.0);
585
586 for op in self.operations.iter().rev() {
588 match op {
589 Operation::Add {
590 result,
591 left,
592 right,
593 } => {
594 if let Some(&grad) = gradients.get(result) {
595 *gradients.entry(left.clone()).or_insert(0.0) += grad;
596 *gradients.entry(right.clone()).or_insert(0.0) += grad;
597 }
598 }
599 Operation::Mul {
600 result,
601 left,
602 right,
603 } => {
604 if let Some(&grad) = gradients.get(result) {
605 let left_val = self.variables[left];
606 let right_val = self.variables[right];
607 *gradients.entry(left.clone()).or_insert(0.0) += grad * right_val;
608 *gradients.entry(right.clone()).or_insert(0.0) += grad * left_val;
609 }
610 }
611 _ => {}
612 }
613 }
614
615 inputs
617 .iter()
618 .map(|&input| {
619 (
620 input.to_string(),
621 gradients.get(input).copied().unwrap_or(0.0),
622 )
623 })
624 .collect()
625 }
626}
627
628pub mod optimizers {
630 use super::*;
631
632 pub trait Optimizer {
634 fn step(&mut self, params: &mut HashMap<String, f64>, gradients: &HashMap<String, f64>);
636
637 fn reset(&mut self);
639 }
640
641 pub struct SGD {
643 learning_rate: f64,
644 momentum: f64,
645 velocities: HashMap<String, f64>,
646 }
647
648 impl SGD {
649 pub fn new(learning_rate: f64, momentum: f64) -> Self {
650 Self {
651 learning_rate,
652 momentum,
653 velocities: HashMap::new(),
654 }
655 }
656 }
657
658 impl Optimizer for SGD {
659 fn step(&mut self, params: &mut HashMap<String, f64>, gradients: &HashMap<String, f64>) {
660 for (name, grad) in gradients {
661 let velocity = self.velocities.entry(name.clone()).or_insert(0.0);
662 *velocity = self.momentum * *velocity - self.learning_rate * grad;
663
664 if let Some(param) = params.get_mut(name) {
665 *param += *velocity;
666 }
667 }
668 }
669
670 fn reset(&mut self) {
671 self.velocities.clear();
672 }
673 }
674
675 pub struct Adam {
677 learning_rate: f64,
678 beta1: f64,
679 beta2: f64,
680 epsilon: f64,
681 t: usize,
682 m: HashMap<String, f64>,
683 v: HashMap<String, f64>,
684 }
685
686 impl Adam {
687 pub fn new(learning_rate: f64) -> Self {
688 Self {
689 learning_rate,
690 beta1: 0.9,
691 beta2: 0.999,
692 epsilon: 1e-8,
693 t: 0,
694 m: HashMap::new(),
695 v: HashMap::new(),
696 }
697 }
698 }
699
700 impl Optimizer for Adam {
701 fn step(&mut self, params: &mut HashMap<String, f64>, gradients: &HashMap<String, f64>) {
702 self.t += 1;
703 let t = self.t as f64;
704
705 for (name, grad) in gradients {
706 let m_t = self.m.entry(name.clone()).or_insert(0.0);
707 let v_t = self.v.entry(name.clone()).or_insert(0.0);
708
709 *m_t = self.beta1 * *m_t + (1.0 - self.beta1) * grad;
711 *v_t = self.beta2 * *v_t + (1.0 - self.beta2) * grad * grad;
712
713 let m_hat = *m_t / (1.0 - self.beta1.powf(t));
715 let v_hat = *v_t / (1.0 - self.beta2.powf(t));
716
717 if let Some(param) = params.get_mut(name) {
719 *param -= self.learning_rate * m_hat / (v_hat.sqrt() + self.epsilon);
720 }
721 }
722 }
723
724 fn reset(&mut self) {
725 self.t = 0;
726 self.m.clear();
727 self.v.clear();
728 }
729 }
730
731 pub struct QNG {
733 learning_rate: f64,
734 regularization: f64,
735 }
736
737 impl QNG {
738 pub fn new(learning_rate: f64, regularization: f64) -> Self {
739 Self {
740 learning_rate,
741 regularization,
742 }
743 }
744 }
745
746 impl Optimizer for QNG {
747 fn step(&mut self, params: &mut HashMap<String, f64>, gradients: &HashMap<String, f64>) {
748 for (name, grad) in gradients {
750 if let Some(param) = params.get_mut(name) {
751 *param -= self.learning_rate * grad;
752 }
753 }
754 }
755
756 fn reset(&mut self) {}
757 }
758}
759
760#[cfg(test)]
761mod tests {
762 use super::*;
763
764 #[test]
765 fn test_autodiff_basic() {
766 let mut autodiff = AutoDiff::new();
767
768 autodiff.register_parameter(DifferentiableParam::new("x", 2.0));
770 autodiff.register_parameter(DifferentiableParam::new("y", 3.0));
771
772 let graph = ComputationNode::Mul(
774 Box::new(ComputationNode::Parameter("x".to_string())),
775 Box::new(ComputationNode::Parameter("y".to_string())),
776 );
777 autodiff.set_graph(graph);
778
779 let result = autodiff.forward().expect("forward pass should succeed");
781 assert_eq!(result, 6.0);
782
783 autodiff
785 .backward(1.0)
786 .expect("backward pass should succeed");
787 let gradients = autodiff.gradients();
788
789 assert_eq!(gradients["x"], 3.0); assert_eq!(gradients["y"], 2.0); }
792
793 #[test]
794 fn test_gradient_tape() {
795 let mut tape = GradientTape::new();
796
797 let x = tape.variable("x", 2.0);
798 let y = tape.variable("y", 3.0);
799 let z = tape.mul(&x, &y);
800
801 let gradients = tape.gradient(&z, &[&x, &y]);
802
803 assert_eq!(gradients[&x], 3.0);
804 assert_eq!(gradients[&y], 2.0);
805 }
806
807 #[test]
808 fn test_optimizers() {
809 use optimizers::*;
810
811 let mut params = HashMap::new();
812 params.insert("x".to_string(), 5.0);
813
814 let mut gradients = HashMap::new();
815 gradients.insert("x".to_string(), 2.0);
816
817 let mut sgd = SGD::new(0.1, 0.0);
819 sgd.step(&mut params, &gradients);
820 assert!((params["x"] - 4.8).abs() < 1e-6);
821
822 params.insert("x".to_string(), 5.0);
824 let mut adam = Adam::new(0.1);
825 adam.step(&mut params, &gradients);
826 assert!(params["x"] < 5.0); }
828
829 #[test]
830 fn test_parameter_shift() {
831 let executor = |params: &[f64]| -> f64 { params[0].cos() + params[1].sin() };
832
833 let qad = QuantumAutoDiff::new(executor);
834 let params = vec![PI / 4.0, PI / 3.0];
835
836 let gradients = qad
837 .parameter_shift_gradients(¶ms, PI / 2.0)
838 .expect("parameter shift gradients should succeed");
839 assert_eq!(gradients.len(), 2);
840 }
841
842 #[test]
843 fn test_expectation_node_without_executor_errors_honestly() {
844 let mut autodiff = AutoDiff::new();
848 autodiff.register_parameter(DifferentiableParam::new("theta", PI / 4.0));
849 autodiff.set_graph(ComputationNode::Expectation {
850 circuit_params: vec!["theta".to_string()],
851 observable: "Z".to_string(),
852 });
853
854 let forward_result = autodiff.forward();
855 assert!(forward_result.is_err());
856 match forward_result {
857 Err(MLError::NotSupported(_)) => {}
858 other => panic!("expected MLError::NotSupported, got {other:?}"),
859 }
860 }
861
862 #[test]
863 fn test_expectation_node_real_parameter_shift_gradient() {
864 let executor = |params: &[f64], _observable: &str| -> f64 { params[0].cos() };
871
872 let mut autodiff = AutoDiff::new().with_executor(executor);
873 let theta = PI / 3.0;
874 autodiff.register_parameter(DifferentiableParam::new("theta", theta));
875 autodiff.set_graph(ComputationNode::Expectation {
876 circuit_params: vec!["theta".to_string()],
877 observable: "Z".to_string(),
878 });
879
880 let forward_value = autodiff.forward().expect("forward should succeed");
881 assert!((forward_value - theta.cos()).abs() < 1e-9);
882
883 autodiff.backward(1.0).expect("backward should succeed");
884 let gradients = autodiff.gradients();
885
886 let expected_gradient = -theta.sin();
887 assert!(
888 (gradients["theta"] - expected_gradient).abs() < 1e-6,
889 "expected d/dtheta cos(theta) = {expected_gradient}, got {}",
890 gradients["theta"]
891 );
892 assert!((gradients["theta"] - 0.5).abs() > 1e-3);
894 }
895}