use crate::common_matrices;
use crate::linalg::Tensor;
use crate::linalg::Trace;
use crate::states::StateData::{Mixed, Pure, Stabilizer};
use crate::tableau::Tableau;
use crate::QubitSized;
use crate::C64;
use core::fmt::Display;
use ndarray::{Array1, Array2, Axis};
use num_traits::One;
use serde::{Deserialize, Serialize};
use std::convert::TryInto;
#[cfg(feature = "python")]
use pyo3::prelude::*;
#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum StateData {
Pure(Array1<C64>),
Mixed(Array2<C64>),
Stabilizer(Tableau),
}
pub type State = QubitSized<StateData>;
impl Display for State {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
write!(
f,
"Register state on {} qubits ({} representation)\nData:\n{}",
self.n_qubits,
match self.data {
Pure(_) => "state vector",
Mixed(_) => "density operator",
StateData::Stabilizer(_) => "stabilizer tableau",
},
self.data
)
}
}
impl Display for StateData {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::result::Result<(), std::fmt::Error> {
match self {
Pure(psi) => write!(f, "{}", psi),
Mixed(rho) => write!(f, "{}", rho),
StateData::Stabilizer(tableau) => write!(f, "{}", tableau),
}
}
}
impl State {
pub fn new_mixed(n_qubits: usize) -> State {
let new_dim = 2usize.pow(n_qubits.try_into().unwrap());
State {
n_qubits,
data: Mixed(common_matrices::elementary_matrix(
(0, 0),
(new_dim, new_dim),
)),
}
}
pub fn new_stabilizer(n_qubits: usize) -> State {
State {
n_qubits,
data: Stabilizer(Tableau::new(n_qubits)),
}
}
pub fn new_pure(n_qubits: usize) -> State {
let new_dim = 2usize.pow(n_qubits.try_into().unwrap());
State {
n_qubits,
data: Pure(common_matrices::elementary_vec(0, new_dim)),
}
}
pub fn as_json(&self) -> String {
serde_json::to_string(&self).unwrap()
}
pub fn extend(&self, n_qubits: usize) -> State {
let new_dim = 2usize.pow(n_qubits.try_into().unwrap());
State {
n_qubits: self.n_qubits + n_qubits,
data: match &self.data {
Pure(psi) => Pure(psi.tensor(&common_matrices::elementary_vec(0, new_dim))),
Mixed(rho) => Mixed(rho.tensor(&common_matrices::elementary_matrix(
(0, 0),
(new_dim, new_dim),
))),
_ => todo!(),
},
}
}
pub fn to_mixed(&self) -> State {
State {
n_qubits: self.n_qubits,
data: match &self.data {
Mixed(rho) => Mixed(rho.clone()),
Pure(psi) => Mixed({
let psi = psi.view().insert_axis(Axis(1));
psi.t().map(|e| e.conj()) * psi
}),
_ => todo!(),
},
}
}
pub fn get_tableau(&self) -> Option<&Tableau> {
match self.data {
Stabilizer(ref tableau) => Some(tableau),
_ => None,
}
}
}
impl Trace for &State {
type Output = C64;
fn trace(self) -> Self::Output {
match &self.data {
Pure(_) | StateData::Stabilizer(_) => C64::one(),
Mixed(ref rho) => (&rho).trace(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn trace_pure_is_one() {
let pure = State {
n_qubits: 1usize,
data: Pure(common_matrices::elementary_vec(0, 2)),
};
assert_eq!(pure.trace(), C64::one());
}
}