Skip to main content

tnc/tensornetwork/
tensordata.rs

1use std::path::PathBuf;
2
3use approx::AbsDiffEq;
4use ndarray::ArrayD;
5use num_complex::Complex64;
6use serde::{Deserialize, Serialize};
7
8use crate::{
9    gates::{load_gate, load_gate_adjoint, matrix_adjoint_inplace},
10    io::hdf5::load_data,
11};
12
13pub type DataTensor = ArrayD<Complex64>;
14
15/// The data of a tensor.
16#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
17pub enum TensorData {
18    /// No data.
19    #[default]
20    None,
21    /// The data is loaded from a HDF5 file.
22    File((PathBuf, bool)),
23    /// A quantum gate. The name must be registered in the gates module.
24    Gate((String, Vec<f64>, bool)),
25    /// A raw vec of complex numbers.
26    Matrix(DataTensor),
27}
28
29impl TensorData {
30    /// Creates a new tensor from raw (flat) data.
31    #[must_use]
32    pub fn new_from_data(dimensions: &[usize], data: Vec<Complex64>) -> Self {
33        Self::Matrix(ArrayD::from_shape_vec(dimensions, data).unwrap())
34    }
35
36    /// Consumes the tensor data and returns the contained tensor.
37    pub fn into_data(self) -> DataTensor {
38        match self {
39            TensorData::None => panic!("Cannot convert uncontracted tensor to data"),
40            TensorData::File((filename, adjoint)) => {
41                let mut data = load_data(filename).unwrap();
42                if adjoint {
43                    matrix_adjoint_inplace(&mut data);
44                }
45                data
46            }
47            TensorData::Gate((gatename, angles, adjoint)) => {
48                if adjoint {
49                    load_gate_adjoint(&gatename, &angles)
50                } else {
51                    load_gate(&gatename, &angles)
52                }
53            }
54            TensorData::Matrix(tensor) => tensor,
55        }
56    }
57
58    /// Returns the adjoint of this data.
59    pub fn adjoint(self) -> Self {
60        match self {
61            TensorData::None => TensorData::None,
62            TensorData::File((filename, adjoint)) => TensorData::File((filename, !adjoint)),
63            TensorData::Gate((name, params, adjoint)) => TensorData::Gate((name, params, !adjoint)),
64            TensorData::Matrix(mut tensor) => {
65                matrix_adjoint_inplace(&mut tensor);
66                TensorData::Matrix(tensor)
67            }
68        }
69    }
70}
71
72impl AbsDiffEq for TensorData {
73    type Epsilon = f64;
74
75    fn default_epsilon() -> Self::Epsilon {
76        f64::EPSILON
77    }
78
79    fn abs_diff_eq(&self, other: &Self, epsilon: Self::Epsilon) -> bool {
80        match (self, other) {
81            (TensorData::File(l0), TensorData::File(r0)) => l0 == r0,
82            (
83                TensorData::Gate((name_l, angles_l, adjoint_l)),
84                TensorData::Gate((name_r, angles_r, adjoint_r)),
85            ) => {
86                name_l == name_r
87                    && adjoint_l == adjoint_r
88                    && angles_l
89                        .iter()
90                        .zip(angles_r)
91                        .all(|(l, r)| f64::abs_diff_eq(l, r, epsilon))
92            }
93            (TensorData::Matrix(l0), TensorData::Matrix(r0)) => {
94                DataTensor::abs_diff_eq(l0, r0, epsilon)
95            }
96            (TensorData::None, TensorData::None) => true,
97            _ => false,
98        }
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use approx::assert_abs_diff_eq;
105
106    use super::*;
107
108    #[test]
109    #[should_panic(expected = "assert_abs_diff_eq!")]
110    fn gates_eq_different_name() {
111        let g1 = TensorData::Gate((String::from("cx"), vec![], false));
112        let g2 = TensorData::Gate((String::from("CX"), vec![], false));
113        assert_abs_diff_eq!(&g1, &g2);
114    }
115
116    #[test]
117    #[should_panic(expected = "assert_abs_diff_eq!")]
118    fn gates_eq_adjoint() {
119        let g1 = TensorData::Gate((String::from("h"), vec![], false));
120        let g2 = TensorData::Gate((String::from("h"), vec![], true));
121        assert_abs_diff_eq!(&g1, &g2);
122    }
123
124    #[test]
125    #[should_panic(expected = "assert_abs_diff_eq!")]
126    fn gates_eq_different_angles() {
127        let g1 = TensorData::Gate((String::from("u"), vec![1.4, 2.0, -3.0], false));
128        let g2 = TensorData::Gate((String::from("u"), vec![1.4, -2.0, -3.0], false));
129        assert_abs_diff_eq!(&g1, &g2);
130    }
131
132    #[test]
133    #[should_panic(expected = "assert_abs_diff_eq!")]
134    fn eq_different_data() {
135        let g1 = TensorData::Gate((String::from("u"), vec![1.4, 2.0, -3.0], false));
136        let g2 = TensorData::new_from_data(&[], vec![Complex64::ONE]);
137        assert_abs_diff_eq!(&g1, &g2);
138    }
139}