refraction/src/riemann.rs

use crate::mathx::{make_mat3, Decomp3};
use glam::*;

pub type Tens3 = [Mat3; 3];

pub trait Metric {
	fn sqrt_at(&self, pos: Vec3) -> Decomp3;

	fn at(&self, pos: Vec3) -> Mat3 {
		self.sqrt_at(pos).square().into()
	}

	fn inverse_at(&self, pos: Vec3) -> Mat3 {
		self.sqrt_at(pos).square().inverse().into()
	}

	fn part_derivs_at(&self, pos: Vec3) -> Tens3 {
		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: Vec3, v: Vec3) -> f32 {
		v.dot(self.at(at) * v).sqrt()
	}

	fn normalize_vec_at(&self, at: Vec3, v: Vec3) -> Vec3 {
		v / self.vec_length_at(at, v)
	}
}

pub struct TraceIter<'a, M: Metric> {
	space: &'a M,
	p: Vec3,
	v: Vec3,
}

impl<'a, M: Metric> Iterator for TraceIter<'a, M> {
	type Item = Vec3;

	fn next(&mut self) -> Option<Self::Item> {
		let a: Vec3 = -contract2(krist(self.space, self.p), self.v);
		self.v += a;
		self.p += self.v;
		Some(self.p)
	}
}

pub fn trace_iter<M: Metric>(space: &M, base: Vec3, dir: Vec3, dt: f32) -> TraceIter<M> {
	TraceIter {
		space,
		p: base,
		v: dt * space.normalize_vec_at(base, dir),
	}
}

pub fn krist(space: &impl Metric, pos: Vec3) -> Tens3 {
	// Γ^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_tens3(|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(Vec3) -> Mat3, pos: Vec3, delta: Vec3) -> Mat3 {
	(f(pos + delta) - f(pos - delta)) / (2.0 * delta.length())
}

fn part_deriv(f: impl Fn(Vec3) -> Mat3, pos: Vec3, eps: f32) -> Tens3 {
	[
		dir_deriv(&f, pos, vec3(eps, 0.0, 0.0)),
		dir_deriv(&f, pos, vec3(0.0, eps, 0.0)),
		dir_deriv(&f, pos, vec3(0.0, 0.0, eps)),
	]
}

/// Сворачивает тензор t с вектором u
pub fn contract(t: Tens3, u: Vec3) -> Mat3 {
	mat3(t[0] * u, t[1] * u, t[2] * u).transpose()
}

/// Сворачивает тензор t с вектором v дважды, по второму и третьему индексам.
pub fn contract2(t: Tens3, v: Vec3) -> Vec3 {
	contract(t, v) * v
}

fn make_tens3(f: impl Fn(usize, usize, usize) -> f32) -> Tens3 {
	std::array::from_fn(|i| make_mat3(|j, k| f(i, j, k)))
}

#[test]
fn m3() {
	let m = make_mat3(|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 t3() {
	let t = make_tens3(|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);
}

pub mod samples {
	use glam::{Mat3, Vec3};

	use super::{Decomp3, Metric};

	pub struct ScaledMetric {
		/// Specifies unit size in each cardinal direction. E.g. with scale=(2, 3), vector (1, 0) has length 2 while a unit vector with the same direction is (1/2, 0).
		pub scale: Vec3,
	}

	impl Metric for ScaledMetric {
		fn sqrt_at(&self, _pos: Vec3) -> Decomp3 {
			Decomp3 {
				diag: self.scale,
				ortho: Mat3::IDENTITY,
			}
		}
	}
}

#[cfg(test)]
mod tests {
	use super::*;
	use approx::assert_abs_diff_eq;

	use glam::{vec3, Mat3};
	use rand::{Rng, SeedableRng};

	#[test]
	fn uniform_scaled_metric() {
		let mut rng = rand_pcg::Pcg64Mcg::seed_from_u64(17);
		let metric = samples::ScaledMetric {
			scale: vec3(3., 4., 5.),
		};
		assert_eq!(
			metric.sqrt_at(rng.gen()),
			Decomp3 {
				ortho: Mat3::IDENTITY,
				diag: vec3(3., 4., 5.)
			}
		);
		assert_eq!(
			metric.at(rng.gen()),
			Mat3::from_cols_array(&[9., 0., 0., 0., 16., 0., 0., 0., 25.])
		);
		assert_eq!(
			metric.inverse_at(rng.gen()),
			Mat3::from_cols_array(&[1. / 9., 0., 0., 0., 1. / 16., 0., 0., 0., 1. / 25.])
		);
		assert_eq!(
			metric.part_derivs_at(rng.gen()),
			[Mat3::ZERO, Mat3::ZERO, Mat3::ZERO]
		);
		assert_eq!(metric.vec_length_at(rng.gen(), vec3(1., 0., 0.)), 3.);
		assert_eq!(metric.vec_length_at(rng.gen(), vec3(0., 1., 0.)), 4.);
		assert_eq!(metric.vec_length_at(rng.gen(), vec3(0., 0., 1.)), 5.);
		assert_eq!(metric.vec_length_at(rng.gen(), vec3(1., 1., 0.)), 5.);
		assert_eq!(
			metric.normalize_vec_at(rng.gen(), vec3(1., 0., 0.)),
			vec3(1. / 3., 0., 0.)
		);
		assert_eq!(
			metric.normalize_vec_at(rng.gen(), vec3(0., 1., 0.)),
			vec3(0., 1. / 4., 0.)
		);
		assert_eq!(
			metric.normalize_vec_at(rng.gen(), vec3(1., 1., 0.)),
			vec3(1. / 5., 1. / 5., 0.)
		);
	}

	#[test]
	fn test_trace_iter() {
		let metric = samples::ScaledMetric {
			scale: vec3(2., 4., 3.),
		};
		assert_eq!(
			trace_iter(&metric, vec3(3., 5., 0.), vec3(1., 0., 0.), 1.)
				.nth(7)
				.unwrap(),
			vec3(7., 5., 0.)
		);
		assert_eq!(
			trace_iter(&metric, vec3(3., 5., 0.), vec3(2., 0., 0.), 1.)
				.nth(7)
				.unwrap(),
			vec3(7., 5., 0.)
		);
		assert_eq!(
			trace_iter(&metric, vec3(3., 5., 0.), vec3(1., 0., 0.), 0.5)
				.nth(7)
				.unwrap(),
			vec3(5., 5., 0.)
		);
		assert_eq!(
			trace_iter(&metric, vec3(3., 5., 0.), vec3(0., 1., 0.), 1.)
				.nth(9)
				.unwrap(),
			vec3(3., 7.5, 0.)
		);
		assert_eq!(
			trace_iter(&metric, vec3(3., 5., 0.), vec3(0., 4., 0.), 1.)
				.nth(9)
				.unwrap(),
			vec3(3., 7.5, 0.)
		);
		assert_eq!(
			trace_iter(&metric, vec3(3., 5., 0.), vec3(0., 1., 0.), 0.5)
				.nth(9)
				.unwrap(),
			vec3(3., 6.25, 0.)
		);
		assert_abs_diff_eq!(
			trace_iter(
				&metric,
				vec3(3., 5., 0.),
				vec3(0.5, 0.25, 0.),
				std::f32::consts::SQRT_2
			)
			.nth(7)
			.unwrap(),
			vec3(7., 7., 0.),
			epsilon = 1e-5
		);
	}
}