aeon_tensor/
lib.rs

1//! Crate for manipulating tensors and tensorial quantaties in Rust.
2
3extern crate self as aeon_tensor;
4
5use std::fmt::Debug;
6use std::{marker::PhantomData, ops};
7
8mod compound;
9mod indices;
10pub mod metric;
11mod storage;
12
13pub use compound::{Compound, CompoundIndices, SymSym, SymVec, VecSym, VecSymVec};
14pub use indices::{Gen, GenIndices, Sym, SymIndices, TensorIndex};
15pub use storage::{TensorStorageMut, TensorStorageOwned, TensorStorageRef};
16
17/// Basic tensor object. Depends on dimension (`N`), rank (`R`), Symmetries
18/// (`I`) and storage array (`S`).
19pub struct Tensor<const N: usize, const R: usize, I, S> {
20    /// Internal storage for tensor, simply wraps around `S`.
21    storage: S,
22    _marker: PhantomData<I>,
23}
24
25impl<const N: usize, const R: usize, I, S> Tensor<N, R, I, S> {
26    /// Retrieves the dimension of the tensor.
27    pub fn dim() -> usize {
28        N
29    }
30
31    /// Retrieves the rank of the tensor.
32    pub fn rank() -> usize {
33        R
34    }
35}
36
37impl<const N: usize, const R: usize, I: TensorIndex<N, R>, S: TensorStorageOwned + Default>
38    Tensor<N, R, I, S>
39{
40    /// Constructs a new tensor with undefined internal values.
41    pub fn new() -> Self {
42        let mut storage = S::default();
43        storage.resize(I::count());
44
45        Self {
46            storage,
47            _marker: PhantomData,
48        }
49    }
50
51    /// Constructs a new tensor by repeatidly calling a function on each index.
52    pub fn from_fn(f: impl Fn([usize; R]) -> f64) -> Self {
53        let mut result = Self::new();
54        result.fill_from_fn(f);
55        result
56    }
57
58    /// Constructs a tensor with all components initialized to v.
59    pub fn splat(v: f64) -> Self {
60        let mut result = Self::new();
61        result.fill(v);
62        result
63    }
64
65    /// Constructs a tensor initialized with all components inititialized to zero.
66    pub fn zeros() -> Self {
67        let mut result = Self::new();
68        result.fill(0.0);
69        result
70    }
71
72    /// Constructs a tensor from a tensorial expression.
73    pub fn from_eq<const C: usize>(f: impl Fn([usize; R], [usize; C]) -> f64) -> Self {
74        Self::from_fn(|index| {
75            let mut result = 0.0;
76            <Gen as TensorIndex<N, C>>::for_each_index(|sum| {
77                result += f(index, sum);
78            });
79            result
80        })
81    }
82}
83
84impl<const N: usize, const R: usize, I: TensorIndex<N, R>, S: TensorStorageRef> From<S>
85    for Tensor<N, R, I, S>
86{
87    fn from(value: S) -> Self {
88        assert!(value.buffer().len() == I::count());
89
90        Self {
91            storage: value,
92            _marker: PhantomData,
93        }
94    }
95}
96
97impl<const N: usize, const R: usize, I: TensorIndex<N, R>, S: TensorStorageMut> Tensor<N, R, I, S> {
98    /// Sets all free components of the tensor to v.
99    pub fn fill(&mut self, v: f64) {
100        let buffer = self.storage.buffer_mut();
101        buffer[..I::count()].fill(v);
102    }
103
104    /// Sets values of components of the tensor be invoking the given function.
105    pub fn fill_from_fn(&mut self, f: impl Fn([usize; R]) -> f64) {
106        let buffer = self.storage.buffer_mut();
107
108        let mut offset = 0;
109
110        I::for_each_index(|index| {
111            buffer[offset] = f(index);
112            offset += 1;
113        });
114    }
115}
116
117impl<const N: usize, const R: usize, I: TensorIndex<N, R>, S: TensorStorageRef> Tensor<N, R, I, S> {
118    /// Retrieves the component at the given index of the tensor.
119    pub fn get(&self, index: [usize; R]) -> &f64 {
120        let offset = I::offset_from_index(index);
121        let buffer = self.storage.buffer();
122        &buffer[offset]
123    }
124}
125
126impl<const N: usize, const R: usize, I: TensorIndex<N, R>, S: TensorStorageMut> Tensor<N, R, I, S> {
127    /// Retrieves a mutable reference to the degree of freedom corresponding to the given index
128    /// of the tensor.
129    pub fn get_mut(&mut self, index: [usize; R]) -> &mut f64 {
130        let offset = I::offset_from_index(index);
131        let buffer = self.storage.buffer_mut();
132        &mut buffer[offset]
133    }
134}
135
136impl<const N: usize, const R: usize, I, S: Clone> Clone for Tensor<N, R, I, S> {
137    fn clone(&self) -> Self {
138        Self {
139            storage: self.storage.clone(),
140            _marker: self._marker.clone(),
141        }
142    }
143}
144
145impl<const N: usize, const R: usize, I, S: Copy> Copy for Tensor<N, R, I, S> {}
146
147impl<const N: usize, const R: usize, I, S: Debug> Debug for Tensor<N, R, I, S> {
148    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
149        self.storage.fmt(f)
150    }
151}
152
153impl<const N: usize, const R: usize, I: TensorIndex<N, R>, S: TensorStorageOwned + Default> Default
154    for Tensor<N, R, I, S>
155{
156    fn default() -> Self {
157        Self::new()
158    }
159}
160
161impl<const N: usize, const R: usize, I: TensorIndex<N, R>, S: TensorStorageRef>
162    ops::Index<[usize; R]> for Tensor<N, R, I, S>
163{
164    type Output = f64;
165
166    fn index(&self, index: [usize; R]) -> &Self::Output {
167        self.get(index)
168    }
169}
170
171impl<const N: usize, const R: usize, I: TensorIndex<N, R>, S: TensorStorageMut>
172    ops::IndexMut<[usize; R]> for Tensor<N, R, I, S>
173{
174    fn index_mut(&mut self, index: [usize; R]) -> &mut Self::Output {
175        self.get_mut(index)
176    }
177}
178
179// *****************************
180// Tests ***********************
181// *****************************
182
183#[cfg(test)]
184mod tests {
185    use crate::{Gen, Sym, SymSym, SymVec, Tensor, TensorIndex, VecSym, VecSymVec};
186
187    fn test_index_axioms<const N: usize, const R: usize, I: TensorIndex<N, R>>() {
188        assert_eq!(
189            I::count(),
190            I::indices().count(),
191            "I::count length ({}) doesn't match I::indices().count() ({})",
192            I::count(),
193            I::indices().count()
194        );
195
196        let mut offset = 0;
197        let mut indices = I::indices();
198
199        I::for_each_index(|i| {
200            let Some(j) = indices.next() else {
201                panic!("I::indices length doesn't match I::for_each_index");
202            };
203
204            assert_eq!(
205                i, j,
206                "I::for_each_index doesn't iterate in the same order as I::indices"
207            );
208
209            assert_eq!(
210                offset,
211                I::offset_from_index(i),
212                "I::offset_from_index doesn't iterate in same order as I::indices"
213            );
214
215            offset += 1;
216        });
217
218        assert_eq!(
219            None,
220            indices.next(),
221            "I::indices length doesn't match I::for_each_index"
222        );
223    }
224
225    #[test]
226    fn index_axioms() {
227        test_index_axioms::<4, 5, Gen>();
228        test_index_axioms::<4, 2, Sym>();
229        test_index_axioms::<4, 3, VecSym>();
230        test_index_axioms::<4, 3, SymVec>();
231        test_index_axioms::<3, 4, SymSym>();
232        test_index_axioms::<3, 4, VecSymVec>();
233    }
234
235    #[test]
236    fn symmetric() {
237        let tensor =
238            Tensor::<2, 4, SymSym, [f64; 9]>::from([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
239
240        // rr
241        assert_eq!(tensor[[0, 0, 0, 0]], 1.0);
242        assert_eq!(tensor[[0, 0, 0, 1]], 2.0);
243        assert_eq!(tensor[[0, 0, 1, 0]], 2.0);
244        assert_eq!(tensor[[0, 0, 1, 1]], 3.0);
245        // rz
246        assert_eq!(tensor[[0, 1, 0, 0]], 4.0);
247        assert_eq!(tensor[[0, 1, 0, 1]], 5.0);
248        assert_eq!(tensor[[0, 1, 1, 0]], 5.0);
249        assert_eq!(tensor[[0, 1, 1, 1]], 6.0);
250        // zr
251        assert_eq!(tensor[[1, 0, 0, 0]], 4.0);
252        assert_eq!(tensor[[1, 0, 0, 1]], 5.0);
253        assert_eq!(tensor[[1, 0, 1, 0]], 5.0);
254        assert_eq!(tensor[[1, 0, 1, 1]], 6.0);
255        // zz
256        assert_eq!(tensor[[1, 1, 0, 0]], 7.0);
257        assert_eq!(tensor[[1, 1, 0, 1]], 8.0);
258        assert_eq!(tensor[[1, 1, 1, 0]], 8.0);
259        assert_eq!(tensor[[1, 1, 1, 1]], 9.0);
260
261        let mut indices = <SymSym as TensorIndex<2, 4>>::indices();
262
263        assert_eq!(indices.next(), Some([0, 0, 0, 0]));
264        assert_eq!(indices.next(), Some([0, 0, 1, 0]));
265        assert_eq!(indices.next(), Some([0, 0, 1, 1]));
266        assert_eq!(indices.next(), Some([1, 0, 0, 0]));
267        assert_eq!(indices.next(), Some([1, 0, 1, 0]));
268        assert_eq!(indices.next(), Some([1, 0, 1, 1]));
269        assert_eq!(indices.next(), Some([1, 1, 0, 0]));
270        assert_eq!(indices.next(), Some([1, 1, 1, 0]));
271        assert_eq!(indices.next(), Some([1, 1, 1, 1]));
272        assert_eq!(indices.next(), None);
273    }
274
275    #[test]
276    fn general() {
277        let tensor =
278            Tensor::<3, 2, Gen, [f64; 9]>::from([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]);
279
280        assert_eq!(tensor[[0, 0]], 1.0);
281        assert_eq!(tensor[[0, 1]], 2.0);
282        assert_eq!(tensor[[0, 2]], 3.0);
283        assert_eq!(tensor[[1, 0]], 4.0);
284        assert_eq!(tensor[[1, 1]], 5.0);
285        assert_eq!(tensor[[1, 2]], 6.0);
286        assert_eq!(tensor[[2, 0]], 7.0);
287        assert_eq!(tensor[[2, 1]], 8.0);
288        assert_eq!(tensor[[2, 2]], 9.0);
289    }
290}
aeon_tensor/lib.rs

aeon_tensor/
lib.rs
