149 lines
3.6 KiB
Rust
149 lines
3.6 KiB
Rust
use glam::*;
|
||
|
||
pub struct Decomp2 {
|
||
pub ortho: Mat2,
|
||
pub diag: Vec2,
|
||
}
|
||
|
||
impl Decomp2 {
|
||
fn square(&self) -> Self {
|
||
Self {
|
||
ortho: self.ortho,
|
||
diag: self.diag * self.diag,
|
||
}
|
||
}
|
||
|
||
pub(crate) fn inverse(&self) -> Self {
|
||
Self {
|
||
ortho: self.ortho,
|
||
diag: Vec2::splat(1.0) / self.diag,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl From<Decomp2> for Mat2 {
|
||
fn from(value: Decomp2) -> Self {
|
||
value.ortho.transpose() * Mat2::from_diagonal(value.diag) * value.ortho
|
||
}
|
||
}
|
||
|
||
pub type Tens2 = [Mat2; 2];
|
||
|
||
pub trait Metric {
|
||
fn sqrt_at(&self, pos: Vec2) -> Decomp2;
|
||
|
||
fn at(&self, pos: Vec2) -> Mat2 {
|
||
self.sqrt_at(pos).square().into()
|
||
}
|
||
|
||
fn inverse_at(&self, pos: Vec2) -> Mat2 {
|
||
self.sqrt_at(pos).square().inverse().into()
|
||
}
|
||
|
||
fn part_derivs_at(&self, pos: Vec2) -> Tens2 {
|
||
part_deriv(|p| self.at(p), pos, 1.0 / 1024.0) // division by such eps is exact which is good for overall precision
|
||
}
|
||
|
||
fn vec_length_at(&self, at: Vec2, v: Vec2) -> f32 {
|
||
v.dot(self.at(at) * v).sqrt()
|
||
}
|
||
|
||
fn normalize_vec_at(&self, at: Vec2, v: Vec2) -> Vec2 {
|
||
v / self.vec_length_at(at, v)
|
||
}
|
||
|
||
fn globalize(&self, at: Vec2, v: Vec2) -> Vec2 {
|
||
Mat2::from(self.sqrt_at(at).inverse()) * v
|
||
}
|
||
}
|
||
|
||
pub struct TraceIter<'a, M: Metric> {
|
||
space: &'a M,
|
||
p: Vec2,
|
||
v: Vec2,
|
||
dt: f32,
|
||
}
|
||
|
||
impl<'a, M: Metric> Iterator for TraceIter<'a, M> {
|
||
type Item = Vec2;
|
||
|
||
fn next(&mut self) -> Option<Self::Item> {
|
||
let a: Vec2 = -contract2(krist(self.space, self.p), self.v);
|
||
self.v = self.v + a * self.dt;
|
||
self.p = self.p + self.v * self.dt;
|
||
Some(self.p)
|
||
}
|
||
}
|
||
|
||
pub fn trace_iter<M: Metric>(space: &M, base: Vec2, dir: Vec2, dt: f32) -> TraceIter<M> {
|
||
TraceIter {
|
||
space,
|
||
p: base,
|
||
v: space.normalize_vec_at(base, dir),
|
||
dt,
|
||
}
|
||
}
|
||
|
||
pub fn krist(space: &impl Metric, pos: Vec2) -> Tens2 {
|
||
// Γ^i_k_l = .5 * g^i^m * (g_m_k,l + g_m_l,k - g_k_l,m)
|
||
let g = &space.inverse_at(pos); // с верхними индексами
|
||
let d = space.part_derivs_at(pos);
|
||
// ret[i][l][k] = sum((m) => .5f * g[m][i] * (d[k][l][m] + d[l][k][m] - d[m][k][l]))
|
||
make_tens2(|i, l, k| 0.5 * (0..2).map(|m| g.col(m)[i] * (d[l].col(k)[m] + d[k].col(m)[l] - d[m].col(k)[l])).sum::<f32>())
|
||
}
|
||
|
||
fn dir_deriv(f: impl Fn(Vec2) -> Mat2, pos: Vec2, delta: Vec2) -> Mat2 {
|
||
(f(pos + delta) - f(pos - delta)) / (2.0 * delta.length())
|
||
}
|
||
|
||
fn part_deriv(f: impl Fn(Vec2) -> Mat2, pos: Vec2, eps: f32) -> Tens2 {
|
||
[
|
||
dir_deriv(&f, pos, vec2(eps, 0.0)),
|
||
dir_deriv(&f, pos, vec2(0.0, eps)),
|
||
]
|
||
}
|
||
|
||
/// Сворачивает тензор t с вектором u
|
||
pub fn contract(t: Tens2, u: Vec2) -> Mat2 {
|
||
mat2(t[0] * u, t[1] * u).transpose()
|
||
}
|
||
|
||
/// Сворачивает тензор t с вектором v дважды, по второму и третьему индексам.
|
||
pub fn contract2(t: Tens2, v: Vec2) -> Vec2 {
|
||
contract(t, v) * v
|
||
}
|
||
|
||
fn make_vec2(f: impl Fn(usize) -> f32) -> Vec2 {
|
||
Vec2::from_array(std::array::from_fn(|i| f(i)))
|
||
}
|
||
|
||
fn make_mat2(f: impl Fn(usize, usize) -> f32) -> Mat2 {
|
||
Mat2::from_cols_array_2d(&std::array::from_fn(|i| std::array::from_fn(|j| f(i, j))))
|
||
}
|
||
|
||
fn make_tens2(f: impl Fn(usize, usize, usize) -> f32) -> Tens2 {
|
||
std::array::from_fn(|i| make_mat2(|j, k| f(i, j, k)))
|
||
}
|
||
|
||
#[test]
|
||
fn m2() {
|
||
let m = make_mat2(|i, j| (i + 2 * j) as f32);
|
||
assert_eq!(m.col(0)[0], 0.0);
|
||
assert_eq!(m.col(1)[0], 1.0);
|
||
assert_eq!(m.col(0)[1], 2.0);
|
||
assert_eq!(m.col(1)[1], 3.0);
|
||
}
|
||
|
||
#[test]
|
||
fn t2() {
|
||
let t = make_tens2(|i, j, k| (i + 2 * j + 4 * k) as f32);
|
||
assert_eq!(t[0].col(0)[0], 0.0);
|
||
assert_eq!(t[1].col(0)[0], 1.0);
|
||
assert_eq!(t[0].col(1)[0], 2.0);
|
||
assert_eq!(t[1].col(1)[0], 3.0);
|
||
assert_eq!(t[0].col(0)[1], 4.0);
|
||
assert_eq!(t[1].col(0)[1], 5.0);
|
||
assert_eq!(t[0].col(1)[1], 6.0);
|
||
assert_eq!(t[1].col(1)[1], 7.0);
|
||
}
|