First quantum calculation

1 files changed, 34 insertions, 5 deletions

diff --git a/libpsi/src/core/component.rs b/libpsi/src/core/component.rs
index bd18b52..51e0551 100644
--- a/libpsi/src/core/component.rs
+++ b/libpsi/src/core/component.rs
@@ -1,7 +1,13 @@
 use crate::{ColumnVector, Complex, Matrix, VectorMatrix};
+use std::ops;
 
-pub type QuantumRegister = ColumnVector<Complex<f64>>;
 pub type QuantumBit = ColumnVector<Complex<f64>>;
+pub type QuantumGate = Matrix<Complex<f64>>;
+
+pub struct QuantumRegister {
+    state: ColumnVector<Complex<f64>>,
+    qubits: Vec<QuantumBit>,
+}
 
 #[macro_export]
 macro_rules! count {
@@ -27,8 +33,8 @@ macro_rules! quantum_register {
     ($($bit:expr),*) => {
         {
             const N: usize = count!($($bit),*);
-            let bits: [QuantumBit; N] = [$($bit),*];
-            QuantumRegister::from(&bits)
+            let mut bits: [QuantumBit; N] = [$($bit),*];
+            QuantumRegister::from(&mut bits)
         }
     };
 }
@@ -40,12 +46,35 @@ impl QuantumBit {
 }
 
 impl QuantumRegister {
-    pub fn from(bits: &[QuantumBit]) -> QuantumRegister {
+    pub fn from(bits: &mut [QuantumBit]) -> QuantumRegister {
         let matrices: Vec<Matrix<Complex<f64>>> = bits.iter().map(|bit| bit.to_matrix()).collect();
         let mut result = matrices[0].clone();
         for matrix in &matrices[1..] {
             result = result.kronecker(matrix);
         }
-        ColumnVector::from_matrix(&result)
+
+        QuantumRegister {
+            qubits: bits.to_vec(),
+            state: ColumnVector::from_matrix(&result),
+        }
+    }
+
+    pub fn apply(&mut self, gate: &QuantumGate, index: usize) {
+        let result: ColumnVector<Complex<f64>> = self.state.mul_matrix(gate).unwrap();
+        self.qubits[index] = result;
+    }
+}
+
+impl ops::Index<usize> for QuantumRegister {
+    type Output = QuantumBit;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        &self.qubits[index]
+    }
+}
+
+impl ops::IndexMut<usize> for QuantumRegister {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        &mut self.qubits[index]
     }
 }