aeon_tensor/metric/
mod.rs

1//! Module containing common operations on manifolds with metrics.
2
3use crate::{
4    Gen, Sym, SymSym, SymVec, Tensor, TensorIndex, TensorStorageOwned, TensorStorageRef, VecSym,
5    VecSymVec,
6};
7
8mod dims;
9
10pub use dims::d2;
11
12pub trait Space<const N: usize>: Clone + Copy {
13    type VecStore: TensorStorageOwned + Default + Clone;
14    type MatStore: TensorStorageOwned + Default + Clone;
15    type SymStore: TensorStorageOwned + Default + Clone;
16    type SymVecStore: TensorStorageOwned + Default + Clone;
17    type SymSymStore: TensorStorageOwned + Default + Clone;
18    type SymVecVecStore: TensorStorageOwned + Default + Clone;
19
20    /// Sums over all indices of the given rank.
21    fn sum<const R: usize>(f: impl Fn([usize; R]) -> f64) -> f64 {
22        let mut result = 0.0;
23        <Gen as TensorIndex<N, R>>::for_each_index(|idx| result += f(idx));
24        result
25    }
26
27    /// Constructs an `N` dimensional vector.
28    fn vector(f: impl Fn([usize; 1]) -> f64) -> Tensor<N, 1, Gen, Self::VecStore> {
29        Tensor::from_fn(f)
30    }
31
32    /// Constructs an `N` dimensional symmetric matrix.
33    fn symmetric(f: impl Fn([usize; 2]) -> f64) -> Tensor<N, 2, Sym, Self::SymStore> {
34        Tensor::from_fn(f)
35    }
36
37    /// Constructs an `N` dimensional general matrix.
38    fn matrix(f: impl Fn([usize; 2]) -> f64) -> Tensor<N, 2, Gen, Self::MatStore> {
39        Tensor::from_fn(f)
40    }
41}
42
43/// Space where all tensors are stored statically and unboxed.
44#[derive(Clone, Copy)]
45pub struct Static;
46
47/// The number of components of an `n` dimensional symmetric matrix.
48const fn sym(n: usize) -> usize {
49    n * (n + 1) / 2
50}
51
52macro_rules! impl_space {
53    ($N:literal) => {
54        impl Space<$N> for Static {
55            type VecStore = [f64; const { $N }];
56            type MatStore = [f64; const { $N * $N }];
57            type SymStore = [f64; const { sym($N) }];
58            type SymVecStore = [f64; const { sym($N) * $N }];
59            type SymSymStore = [f64; const { sym($N) * sym($N) }];
60            type SymVecVecStore = [f64; const { sym($N) * $N * $N }];
61        }
62    };
63}
64
65impl_space!(1);
66impl_space!(2);
67
68/// A metric (along with first and second derivatives) defined on a point on a manifold.
69pub struct Metric<const N: usize, S: Space<N>> {
70    pub value: Tensor<N, 2, Sym, S::SymStore>,
71    pub derivs: Tensor<N, 3, SymVec, S::SymVecStore>,
72    pub derivs2: Tensor<N, 4, SymSym, S::SymSymStore>,
73}
74
75impl<const N: usize, S: Space<N>> Metric<N, S> {
76    /// Constructs a new metric from the constituent components and partial derivatives.
77    pub fn new(
78        value: Tensor<N, 2, Sym, S::SymStore>,
79        derivs: Tensor<N, 3, SymVec, S::SymVecStore>,
80        derivs2: Tensor<N, 4, SymSym, S::SymSymStore>,
81    ) -> Self {
82        Self {
83            value,
84            derivs,
85            derivs2,
86        }
87    }
88}
89
90impl<S: Space<2>> Metric<2, S> {
91    /// Computes the determinate of a metric.
92    pub fn det(&self) -> MetricDet<2, S> {
93        let value = self.value[[0, 0]] * self.value[[1, 1]] - self.value[[1, 0]].powi(2);
94        let derivs = Tensor::from_fn(|[a]| {
95            self.derivs[[0, 0, a]] * self.value[[1, 1]]
96                + self.value[[0, 0]] * self.derivs[[1, 1, a]]
97                - 2.0 * self.value[[0, 1]] * self.derivs[[0, 1, a]]
98        });
99
100        MetricDet { value, derivs }
101    }
102
103    /// Computes the inverse of a metric.
104    pub fn inv(&self, det: &MetricDet<2, S>) -> MetricInv<2, S> {
105        let factor = det.value.recip();
106        let factor_derivs: Tensor<2, 1, Gen, S::VecStore> =
107            Tensor::from_fn(|[a]| -factor.powi(2) * det.derivs[[a]]);
108
109        let mut value = Tensor::new();
110        value[[0, 0]] = factor * self.value[[1, 1]];
111        value[[1, 0]] = -factor * self.value[[1, 0]];
112        debug_assert_eq!(value[[1, 0]], value[[0, 1]]);
113        value[[1, 1]] = factor * self.value[[0, 0]];
114
115        let mut derivs = Tensor::new();
116
117        for a in 0..2 {
118            derivs[[0, 0, a]] =
119                factor_derivs[[a]] * self.value[[1, 1]] + factor * self.derivs[[1, 1, a]];
120            derivs[[1, 0, a]] =
121                -factor_derivs[[a]] * self.value[[1, 0]] - factor * self.derivs[[1, 0, a]];
122            debug_assert_eq!(derivs[[0, 1, a]], derivs[[1, 0, a]]);
123            derivs[[1, 1, a]] =
124                factor_derivs[[a]] * self.value[[0, 0]] + factor * self.derivs[[0, 0, a]];
125        }
126
127        MetricInv { value, derivs }
128    }
129}
130
131impl<const N: usize, S: Space<N>> Metric<N, S> {
132    // Computes killing's equation 𝓛ₓgₐᵦ for the given vector field X.
133    pub fn killing(&self, vector: &VectorC1<N, S>) -> Tensor<N, 2, Sym, S::SymStore> {
134        SymmetricC1 {
135            value: self.value.clone(),
136            derivs: self.derivs.clone(),
137        }
138        .lie_derivative(vector)
139    }
140}
141
142/// The determinate of a metric, along with partial derivatives.
143pub struct MetricDet<const N: usize, S: Space<N>> {
144    pub value: f64,
145    pub derivs: Tensor<N, 1, Gen, S::VecStore>,
146}
147
148/// The inverse of a metric, along with partial derivatives.
149pub struct MetricInv<const N: usize, S: Space<N>> {
150    pub value: Tensor<N, 2, Sym, S::SymStore>,
151    pub derivs: Tensor<N, 3, SymVec, S::SymVecStore>,
152}
153
154impl<const N: usize, S: Space<N>> MetricInv<N, S> {
155    /// Computes the trace of a fully covariant 2-tensor.
156    pub fn cotrace<I: TensorIndex<N, 2>, St: TensorStorageRef>(
157        &self,
158        matrix: &Tensor<N, 2, I, St>,
159    ) -> f64 {
160        S::sum(|[a, b]| self.value[[a, b]] * matrix[[a, b]])
161    }
162
163    /// Raises the first index of a general r-tensor.
164    pub fn raise_first<const R: usize, I: TensorIndex<N, R>, St: TensorStorageOwned + Default>(
165        &self,
166        tensor: &Tensor<N, R, I, St>,
167    ) -> Tensor<N, R, I, St> {
168        const {
169            if R == 0 {
170                panic!("R must be > 0");
171            }
172        }
173
174        Tensor::from_fn(|idx| {
175            S::sum(|[a]| {
176                let mut tidx = idx;
177                tidx[0] = a;
178                self.value[[idx[0], a]] * tensor[tidx]
179            })
180        })
181    }
182
183    /// Raises the last index of a general r-tensor.
184    pub fn raise_last<const R: usize, I: TensorIndex<N, R>, St: TensorStorageOwned + Default>(
185        &self,
186        tensor: &Tensor<N, R, I, St>,
187    ) -> Tensor<N, R, I, St> {
188        const {
189            if R == 0 {
190                panic!("R must be > 0");
191            }
192        }
193
194        Tensor::from_fn(|idx| {
195            S::sum(|[a]| {
196                let mut tidx = idx;
197                tidx[R - 1] = a;
198                self.value[[idx[R - 1], a]] * tensor[tidx]
199            })
200        })
201    }
202}
203
204/// Christoffel connection symbols and their derivatives defined on
205/// a general metric.
206pub struct ChristoffelSymbol<const N: usize, S: Space<N>> {
207    pub first_kind: Tensor<N, 3, VecSym, S::SymVecStore>,
208    pub first_kind_derivs: Tensor<N, 4, VecSymVec, S::SymVecVecStore>,
209    pub second_kind: Tensor<N, 3, VecSym, S::SymVecStore>,
210    pub second_kind_derivs: Tensor<N, 4, VecSymVec, S::SymVecVecStore>,
211}
212
213impl<const N: usize, S: Space<N>> ChristoffelSymbol<N, S> {
214    /// Computes the ricci tensor from the christoffel symbols.
215    pub fn ricci(&self) -> Tensor<N, 2, Sym, S::SymStore> {
216        Tensor::from_eq(|[i, j], [a]| {
217            let term1: f64 =
218                self.second_kind_derivs[[a, i, j, a]] - self.second_kind_derivs[[a, a, i, j]];
219            let term2 = S::sum(|[b]| {
220                self.second_kind[[a, a, b]] * self.second_kind[[b, i, j]]
221                    - self.second_kind[[a, i, b]] * self.second_kind[[b, a, j]]
222            });
223
224            term1 + term2
225        })
226    }
227}
228
229impl<const N: usize, S: Space<N>> Metric<N, S> {
230    /// Computes Christoffel_symbols for a given metric.
231    pub fn christoffel_symbol(&self, inv: &MetricInv<N, S>) -> ChristoffelSymbol<N, S> {
232        let first_kind = Tensor::from_fn(|[a, b, c]| {
233            0.5 * (self.derivs[[a, c, b]] + self.derivs[[b, a, c]] - self.derivs[[b, c, a]])
234        });
235        let first_kind_derivs = Tensor::from_fn(|[a, b, c, d]| {
236            0.5 * (self.derivs2[[a, c, b, d]] + self.derivs2[[b, a, c, d]]
237                - self.derivs2[[b, c, a, d]])
238        });
239
240        let second_kind =
241            Tensor::from_eq(|[a, b, c], [m]| inv.value[[a, m]] * first_kind[[m, b, c]]);
242
243        let second_kind_derivs = Tensor::from_eq(|[a, b, c, d], [m]| {
244            inv.derivs[[a, m, d]] * first_kind[[m, b, c]]
245                + inv.value[[a, m]] * first_kind_derivs[[m, b, c, d]]
246        });
247
248        ChristoffelSymbol {
249            first_kind,
250            first_kind_derivs,
251            second_kind,
252            second_kind_derivs,
253        }
254    }
255}
256
257/// A C1 scalar field at a single point.
258#[derive(Debug)]
259pub struct ScalarC1<const N: usize, S: Space<N>> {
260    pub value: f64,
261    pub derivs: Tensor<N, 1, Gen, S::VecStore>,
262}
263
264impl<const N: usize, S: Space<N>> ScalarC1<N, S> {
265    pub fn gradient(&self, _connect: &ChristoffelSymbol<N, S>) -> Tensor<N, 1, Gen, S::VecStore> {
266        self.derivs.clone()
267    }
268
269    pub fn lie_derivative(&self, flow: &VectorC1<N, S>) -> f64 {
270        S::sum(|[a]| flow.value[[a]] * self.derivs[[a]])
271    }
272}
273
274impl<const N: usize, S: Space<N>> From<ScalarC2<N, S>> for ScalarC1<N, S> {
275    fn from(value: ScalarC2<N, S>) -> Self {
276        Self {
277            value: value.value,
278            derivs: value.derivs,
279        }
280    }
281}
282
283impl<const N: usize, S: Space<N>> Default for ScalarC1<N, S> {
284    fn default() -> Self {
285        Self {
286            value: Default::default(),
287            derivs: Default::default(),
288        }
289    }
290}
291
292impl<const N: usize, S: Space<N>> Clone for ScalarC1<N, S> {
293    fn clone(&self) -> Self {
294        Self {
295            value: self.value.clone(),
296            derivs: self.derivs.clone(),
297        }
298    }
299}
300
301/// A C2 scalar field at a single point.
302pub struct ScalarC2<const N: usize, S: Space<N>> {
303    pub value: f64,
304    pub derivs: Tensor<N, 1, Gen, S::VecStore>,
305    pub derivs2: Tensor<N, 2, Sym, S::SymStore>,
306}
307
308impl<const N: usize, S: Space<N>> ScalarC2<N, S> {
309    pub fn gradient(&self, _connect: &ChristoffelSymbol<N, S>) -> Tensor<N, 1, Gen, S::VecStore> {
310        self.derivs.clone()
311    }
312
313    pub fn lie_derivative(&self, flow: &VectorC1<N, S>) -> f64 {
314        S::sum(|[a]| flow.value[[a]] * self.derivs[[a]])
315    }
316
317    pub fn hessian(&self, connect: &ChristoffelSymbol<N, S>) -> Tensor<N, 2, Sym, S::SymStore> {
318        Tensor::from_fn(|[a, b]| {
319            let term1 = self.derivs2[[a, b]];
320            let term2 = S::sum(|[d]| -connect.second_kind[[d, a, b]] * self.derivs[[d]]);
321            term1 + term2
322        })
323    }
324}
325
326impl<const N: usize, S: Space<N>> Default for ScalarC2<N, S> {
327    fn default() -> Self {
328        Self {
329            value: Default::default(),
330            derivs: Default::default(),
331            derivs2: Default::default(),
332        }
333    }
334}
335
336impl<const N: usize, S: Space<N>> Clone for ScalarC2<N, S> {
337    fn clone(&self) -> Self {
338        Self {
339            value: self.value.clone(),
340            derivs: self.derivs.clone(),
341            derivs2: self.derivs2.clone(),
342        }
343    }
344}
345
346/// A C1 vector field at a single point.
347pub struct VectorC1<const N: usize, S: Space<N>> {
348    pub value: Tensor<N, 1, Gen, S::VecStore>,
349    pub derivs: Tensor<N, 2, Gen, S::MatStore>,
350}
351
352impl<const N: usize, S: Space<N>> VectorC1<N, S> {
353    pub fn gradient(&self, connect: &ChristoffelSymbol<N, S>) -> Tensor<N, 2, Gen, S::MatStore> {
354        Tensor::from_fn(|[a, c]| {
355            let term1 = self.derivs[[a, c]];
356            let term2 = S::sum(|[d]| -connect.second_kind[[d, a, c]] * self.value[[d]]);
357            term1 + term2
358        })
359    }
360
361    pub fn lie_derivative(&self, flow: &VectorC1<N, S>) -> Tensor<N, 1, Gen, S::VecStore> {
362        Tensor::from_fn(|[a]| {
363            S::sum(|[i]| {
364                flow.value[[i]] * self.derivs[[a, i]] + flow.derivs[[i, a]] * self.value[[i]]
365            })
366        })
367    }
368}
369
370impl<const N: usize, S: Space<N>> Default for VectorC1<N, S> {
371    fn default() -> Self {
372        Self {
373            value: Default::default(),
374            derivs: Default::default(),
375        }
376    }
377}
378
379impl<const N: usize, S: Space<N>> Clone for VectorC1<N, S> {
380    fn clone(&self) -> Self {
381        Self {
382            value: self.value.clone(),
383            derivs: self.derivs.clone(),
384        }
385    }
386}
387
388/// A C1 symmetric matrix field at a single point.
389pub struct SymmetricC1<const N: usize, S: Space<N>> {
390    pub value: Tensor<N, 2, Sym, S::SymStore>,
391    pub derivs: Tensor<N, 3, SymVec, S::SymVecStore>,
392}
393
394impl<const N: usize, S: Space<N>> SymmetricC1<N, S> {
395    pub fn gradient(
396        &self,
397        connect: &ChristoffelSymbol<N, S>,
398    ) -> Tensor<N, 3, SymVec, S::SymVecStore> {
399        Tensor::from_fn(|[i, j, k]| {
400            self.derivs[[i, j, k]]
401                - S::sum(|[m]| connect.second_kind[[m, i, k]] * self.value[[m, j]])
402                - S::sum(|[m]| connect.second_kind[[m, j, k]] * self.value[[m, i]])
403        })
404    }
405
406    pub fn lie_derivative(&self, flow: &VectorC1<N, S>) -> Tensor<N, 2, Sym, S::SymStore> {
407        Tensor::from_fn(|[i, j]| {
408            S::sum(|[m]| {
409                flow.value[[m]] * self.derivs[[i, j, m]]
410                    + self.value[[m, j]] * flow.derivs[[m, i]]
411                    + self.value[[i, m]] * flow.derivs[[m, j]]
412            })
413        })
414    }
415}
416
417impl<const N: usize, S: Space<N>> Default for SymmetricC1<N, S> {
418    fn default() -> Self {
419        Self {
420            value: Default::default(),
421            derivs: Default::default(),
422        }
423    }
424}
425
426impl<const N: usize, S: Space<N>> Clone for SymmetricC1<N, S> {
427    fn clone(&self) -> Self {
428        Self {
429            value: self.value.clone(),
430            derivs: self.derivs.clone(),
431        }
432    }
433}
aeon_tensor/metric/mod.rs

aeon_tensor/metric/
mod.rs
