use super::FixedVector;
use super::FixedMatrix;
use super::tensor::Tensor;
use super::compute_matrix::{compute_lu_decompose, ComputeMatrix};
use super::ode::{OdeSystem, rk4_step_compute, state_to_compute};
use super::kernels;
use super::Scalar;
use crate::fixed_point::fixed_num::{Compute, Fixed, FixedNum};
use crate::fixed_point::width::{Storage, WideInt};
use crate::fixed_point::core_types::errors::OverflowDetected;
pub fn differentiation_step<S: Storage, const F: u32>() -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
Fixed::from_raw(S::one().shl_bits(F - differentiation_exponent::<S, F>()))
}
pub trait MetricProvider<T: Scalar> {
fn dimension(&self) -> usize;
fn metric(&self, p: &FixedVector<T>) -> FixedMatrix<T>;
fn metric_inverse(&self, p: &FixedVector<T>) -> Result<FixedMatrix<T>, OverflowDetected> {
T::invert_matrix(&self.metric(p))
}
fn christoffel_closed_form(&self, _p: &FixedVector<T>) -> Option<Tensor<T>> {
None
}
#[doc(hidden)]
fn christoffel_closed_form_compute(&self, _p: &FixedVector<T>) -> Option<ComputeChristoffel<T>> {
None
}
fn scalar_curvature_closed_form(&self, _p: &FixedVector<T>) -> Option<T> {
None
}
}
#[doc(hidden)]
#[derive(Clone, Debug)]
pub struct ComputeChristoffel<T: FixedNum> {
data: Vec<Compute<T>>,
}
fn central_difference<S: Storage, const F: u32>(a: Compute<Fixed<S, F>>, b: Compute<Fixed<S, F>>) -> Result<Compute<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let mut d = checked_sub(a, b)?;
for _ in 1..differentiation_exponent::<S, F>() {
d = Fixed::<S, F>::compute_checked_add(d, d)?;
}
Ok(d)
}
fn differentiation_exponent<S: Storage, const F: u32>() -> u32 {
(F + 1) / 3
}
fn compute_zero<S: Storage, const F: u32>() -> Compute<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
Fixed::<S, F>::upscale_to_compute(Fixed::ZERO.raw())
}
fn compute_one<S: Storage, const F: u32>() -> Compute<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
Fixed::<S, F>::upscale_to_compute(Fixed::one().raw())
}
fn round_tensor<S: Storage, const F: u32>(shape: &[usize], data: &[Compute<Fixed<S, F>>]) -> Result<Tensor<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let values = data.iter()
.map(|&c| Fixed::<S, F>::downscale_to_storage(c).map(Fixed::from_raw))
.collect::<Result<Vec<_>, _>>()?;
Ok(Tensor::<Fixed<S, F>>::from_data(shape, &values))
}
fn metric_inverse_compute<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
) -> Result<Vec<Compute<Fixed<S, F>>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
let given = provider.metric_inverse(p)?;
let given_c: Vec<Compute<Fixed<S, F>>> = (0..n * n)
.map(|e| Fixed::<S, F>::upscale_to_compute(given.get(e / n, e % n).raw()))
.collect();
let wide = compute_lu_decompose(&ComputeMatrix::<Fixed<S, F>>::from_fixed_matrix(&provider.metric(p)))
.and_then(|lu| lu.inverse());
let Ok(wide) = wide else { return Ok(given_c) };
let unit = compute_one_unit::<S, F>();
let mut out = Vec::with_capacity(n * n);
for e in 0..n * n {
let c = wide.get(e / n, e % n);
let diff = checked_sub(c, given_c[e])?;
let diff = if diff < compute_zero::<S, F>() { Fixed::<S, F>::compute_negate(diff) } else { diff };
if diff > unit { return Ok(given_c); }
out.push(c);
}
Ok(out)
}
fn compute_one_unit<S: Storage, const F: u32>() -> Compute<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
Fixed::<S, F>::upscale_to_compute(S::one())
}
fn metric_partial_compute<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
m: usize,
) -> Result<Vec<Compute<Fixed<S, F>>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let h = differentiation_step::<S, F>();
let n = provider.dimension();
let mut p_plus = p.clone();
let mut p_minus = p.clone();
p_plus[m] = p_plus[m] + h;
p_minus[m] = p_minus[m] - h;
let g_plus = provider.metric(&p_plus);
let g_minus = provider.metric(&p_minus);
let mut result = Vec::with_capacity(n * n);
for i in 0..n {
for j in 0..n {
result.push(central_difference(
Fixed::<S, F>::upscale_to_compute(g_plus.get(i, j).raw()),
Fixed::<S, F>::upscale_to_compute(g_minus.get(i, j).raw()),
)?);
}
}
Ok(result)
}
pub fn christoffel<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
) -> Result<Tensor<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
round_tensor(&[n, n, n], &christoffel_compute(provider, p)?)
}
fn christoffel_compute<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
) -> Result<Vec<Compute<Fixed<S, F>>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
if let Some(gamma) = provider.christoffel_closed_form_compute(p) {
return Ok(gamma.data);
}
if let Some(gamma) = provider.christoffel_closed_form(p) {
let mut out = Vec::with_capacity(n * n * n);
for k in 0..n { for i in 0..n { for j in 0..n {
out.push(Fixed::<S, F>::upscale_to_compute(gamma.get(&[k, i, j]).raw()));
} } }
return Ok(out);
}
let g_inv = metric_inverse_compute(provider, p)?;
let dg: Vec<Vec<Compute<Fixed<S, F>>>> = (0..n)
.map(|m| metric_partial_compute(provider, p, m))
.collect::<Result<_, _>>()?;
let mut gamma = Vec::with_capacity(n * n * n);
for k in 0..n {
for i in 0..n {
for j in 0..n {
let mut acc = compute_zero::<S, F>();
for l in 0..n {
let bracket = checked_sub(
Fixed::<S, F>::compute_checked_add(dg[i][j * n + l], dg[j][l * n + i])?,
dg[l][i * n + j],
)?;
acc = Fixed::<S, F>::compute_checked_add(acc, kernels::compute_product::<Fixed<S, F>>(g_inv[k * n + l], bracket)?)?;
}
gamma.push(Fixed::<S, F>::compute_halve(acc));
}
}
}
Ok(gamma)
}
pub fn riemann_curvature<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
) -> Result<Tensor<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
round_tensor(&[n, n, n, n], &riemann_compute(provider, p)?)
}
fn riemann_compute<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
) -> Result<Vec<Compute<Fixed<S, F>>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
let h = differentiation_step::<S, F>();
let at3 = |a: usize, b: usize, c: usize| (a * n + b) * n + c;
let gamma = christoffel_compute(provider, p)?;
let mut dgamma: Vec<Vec<Compute<Fixed<S, F>>>> = Vec::with_capacity(n);
for j in 0..n {
let mut p_plus = p.clone();
let mut p_minus = p.clone();
p_plus[j] = p_plus[j] + h;
p_minus[j] = p_minus[j] - h;
let gamma_plus = christoffel_compute(provider, &p_plus)?;
let gamma_minus = christoffel_compute(provider, &p_minus)?;
dgamma.push(gamma_plus.iter().zip(&gamma_minus)
.map(|(&a, &b)| central_difference(a, b))
.collect::<Result<_, _>>()?);
}
let mut riemann = Vec::with_capacity(n * n * n * n);
for l in 0..n {
for i in 0..n {
for j in 0..n {
for k in 0..n {
let mut acc = checked_sub(dgamma[j][at3(l, i, k)], dgamma[k][at3(l, i, j)])?;
for m in 0..n {
acc = Fixed::<S, F>::compute_checked_add(acc, kernels::compute_product::<Fixed<S, F>>(gamma[at3(l, j, m)], gamma[at3(m, i, k)])?)?;
acc = checked_sub(acc, kernels::compute_product::<Fixed<S, F>>(gamma[at3(l, k, m)], gamma[at3(m, i, j)])?)?;
}
riemann.push(acc);
}
}
}
}
Ok(riemann)
}
pub fn ricci_tensor<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
) -> Result<FixedMatrix<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
let ricci = ricci_compute(provider, p)?;
let mut out = FixedMatrix::<Fixed<S, F>>::new(n, n);
for i in 0..n {
for j in 0..n {
out.set(i, j, Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(ricci[i * n + j])?));
}
}
Ok(out)
}
fn ricci_compute<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
) -> Result<Vec<Compute<Fixed<S, F>>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
let riemann = riemann_compute(provider, p)?;
let at4 = |a: usize, b: usize, c: usize, d: usize| ((a * n + b) * n + c) * n + d;
let mut ricci = Vec::with_capacity(n * n);
for i in 0..n {
for j in 0..n {
let mut acc = compute_zero::<S, F>();
for k in 0..n {
acc = Fixed::<S, F>::compute_checked_add(acc, riemann[at4(k, i, k, j)])?;
}
ricci.push(acc);
}
}
Ok(ricci)
}
pub fn ricci_from_riemann<S: Storage, const F: u32>(riemann: &Tensor<Fixed<S, F>>, n: usize) -> FixedMatrix<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let mut ricci = FixedMatrix::<Fixed<S, F>>::new(n, n);
for i in 0..n {
for j in 0..n {
let mut sum = Fixed::ZERO;
for k in 0..n {
sum = sum + riemann.get(&[k, i, k, j]);
}
ricci.set(i, j, sum);
}
}
ricci
}
pub fn scalar_curvature<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
) -> Result<Fixed<S, F>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if let Some(r) = provider.scalar_curvature_closed_form(p) {
return Ok(r);
}
let n = provider.dimension();
let g_inv = metric_inverse_compute(provider, p)?;
let ricci = ricci_compute(provider, p)?;
let mut acc = compute_zero::<S, F>();
for e in 0..n * n {
acc = Fixed::<S, F>::compute_checked_add(acc, kernels::compute_product::<Fixed<S, F>>(g_inv[e], ricci[e])?)?;
}
Ok(Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(acc)?))
}
pub fn scalar_from_ricci<S: Storage, const F: u32>(g_inv: &FixedMatrix<Fixed<S, F>>, ricci: &FixedMatrix<Fixed<S, F>>) -> Fixed<S, F>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = g_inv.rows();
let mut g_flat = Vec::with_capacity(n * n);
let mut r_flat = Vec::with_capacity(n * n);
for i in 0..n {
for j in 0..n {
g_flat.push(g_inv.get(i, j).raw());
r_flat.push(ricci.get(i, j).raw());
}
}
Fixed::from_raw(kernels::dot_raw::<Fixed<S, F>>(&g_flat, &r_flat))
}
pub fn sectional_curvature<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
p: &FixedVector<Fixed<S, F>>,
u: &FixedVector<Fixed<S, F>>,
v: &FixedVector<Fixed<S, F>>,
) -> Result<Fixed<S, F>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
let g = provider.metric(p);
let riemann = riemann_compute(provider, p)?;
let at4 = |a: usize, b: usize, c: usize, d: usize| ((a * n + b) * n + c) * n + d;
let uc: Vec<Compute<Fixed<S, F>>> = (0..n).map(|i| Fixed::<S, F>::upscale_to_compute(u[i].raw())).collect();
let vc: Vec<Compute<Fixed<S, F>>> = (0..n).map(|i| Fixed::<S, F>::upscale_to_compute(v[i].raw())).collect();
let lower = |x: &[Compute<Fixed<S, F>>]| -> Result<Vec<Compute<Fixed<S, F>>>, OverflowDetected> {
(0..n).map(|i| {
let mut acc = compute_zero::<S, F>();
for j in 0..n {
acc = Fixed::<S, F>::compute_checked_add(acc, kernels::compute_product::<Fixed<S, F>>(Fixed::<S, F>::upscale_to_compute(g.get(i, j).raw()), x[j])?)?;
}
Ok(acc)
}).collect()
};
let gu = lower(&uc)?;
let gv = lower(&vc)?;
let mut numerator = compute_zero::<S, F>();
for s in 0..n {
for i in 0..n {
for j in 0..n {
for k in 0..n {
let r = riemann[at4(s, i, j, k)];
if r == compute_zero::<S, F>() { continue; }
let term = kernels::compute_product::<Fixed<S, F>>(kernels::compute_product::<Fixed<S, F>>(kernels::compute_product::<Fixed<S, F>>(r, vc[i])?, uc[j])?, vc[k])?;
numerator = Fixed::<S, F>::compute_checked_add(numerator, kernels::compute_product::<Fixed<S, F>>(term, gu[s])?)?;
}
}
}
}
let inner = |x: &[Compute<Fixed<S, F>>], gy: &[Compute<Fixed<S, F>>]| -> Result<Compute<Fixed<S, F>>, OverflowDetected> {
let mut acc = compute_zero::<S, F>();
for i in 0..n { acc = Fixed::<S, F>::compute_checked_add(acc, kernels::compute_product::<Fixed<S, F>>(x[i], gy[i])?)?; }
Ok(acc)
};
let uu = inner(&uc, &gu)?;
let vv = inner(&vc, &gv)?;
let uv = inner(&uc, &gv)?;
let denom = checked_sub(kernels::compute_product::<Fixed<S, F>>(uu, vv)?, kernels::compute_product::<Fixed<S, F>>(uv, uv)?)?;
if denom == compute_zero::<S, F>() {
return Err(OverflowDetected::DomainError);
}
Ok(Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(Fixed::<S, F>::compute_checked_divide(numerator, denom)?)?))
}
pub struct EuclideanMetric {
pub dim: usize,
}
impl<S: Storage, const F: u32> MetricProvider<Fixed<S, F>> for EuclideanMetric
where
Fixed<S, F>: FixedNum<Storage = S>,
{
fn dimension(&self) -> usize { self.dim }
fn metric(&self, _p: &FixedVector<Fixed<S, F>>) -> FixedMatrix<Fixed<S, F>> {
FixedMatrix::<Fixed<S, F>>::identity(self.dim)
}
fn metric_inverse(&self, _p: &FixedVector<Fixed<S, F>>) -> Result<FixedMatrix<Fixed<S, F>>, OverflowDetected> {
Ok(FixedMatrix::<Fixed<S, F>>::identity(self.dim))
}
}
pub struct SphereMetric<T> {
pub radius: T,
}
impl<S: Storage, const F: u32> SphereMetric<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
fn radius_sq_compute(&self) -> Compute<Fixed<S, F>> {
let r = Fixed::<S, F>::upscale_to_compute(self.radius.raw());
Fixed::<S, F>::compute_multiply(r, r)
}
}
impl<S: Storage, const F: u32> MetricProvider<Fixed<S, F>> for SphereMetric<Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
fn dimension(&self) -> usize { 2 }
fn metric(&self, p: &FixedVector<Fixed<S, F>>) -> FixedMatrix<Fixed<S, F>> {
let (sin_t, _) = Fixed::<S, F>::sincos_at_compute_tier(Fixed::<S, F>::upscale_to_compute(p[0].raw()));
let r_sq = self.radius_sq_compute();
let g11 = Fixed::<S, F>::compute_multiply(r_sq, Fixed::<S, F>::compute_multiply(sin_t, sin_t));
let z = Fixed::ZERO;
FixedMatrix::<Fixed<S, F>>::from_slice(2, 2, &[
Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(r_sq)), z,
z, Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(g11)),
])
}
fn christoffel_closed_form(&self, p: &FixedVector<Fixed<S, F>>) -> Option<Tensor<Fixed<S, F>>> {
let gamma = self.christoffel_closed_form_compute(p)?;
Some(round_tensor(&[2, 2, 2], &gamma.data).expect("sphere: Christoffel symbol exceeds storage"))
}
#[doc(hidden)]
fn christoffel_closed_form_compute(&self, p: &FixedVector<Fixed<S, F>>) -> Option<ComputeChristoffel<Fixed<S, F>>> {
let (sin_t, cos_t) = Fixed::<S, F>::sincos_at_compute_tier(Fixed::<S, F>::upscale_to_compute(p[0].raw()));
let mut data = vec![compute_zero::<S, F>(); 8];
data[3] = Fixed::<S, F>::compute_negate(Fixed::<S, F>::compute_multiply(sin_t, cos_t));
if sin_t != compute_zero::<S, F>() {
let cot_t = Fixed::<S, F>::compute_divide(cos_t, sin_t).expect("sphere: cot(theta) exceeds the compute tier");
data[5] = cot_t;
data[6] = cot_t;
}
Some(ComputeChristoffel::<Fixed<S, F>> { data })
}
fn scalar_curvature_closed_form(&self, _p: &FixedVector<Fixed<S, F>>) -> Option<Fixed<S, F>> {
let two = Fixed::<S, F>::compute_add(compute_one::<S, F>(), compute_one::<S, F>());
let r = Fixed::<S, F>::compute_divide(two, self.radius_sq_compute())
.expect("sphere: 2 / r^2 exceeds the compute tier");
Some(Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(r)))
}
}
pub struct HyperbolicMetric;
impl<S: Storage, const F: u32> MetricProvider<Fixed<S, F>> for HyperbolicMetric
where
Fixed<S, F>: FixedNum<Storage = S>,
{
fn dimension(&self) -> usize { 2 }
fn metric(&self, p: &FixedVector<Fixed<S, F>>) -> FixedMatrix<Fixed<S, F>> {
let y = Fixed::<S, F>::upscale_to_compute(p[1].raw());
let scale = Fixed::<S, F>::compute_divide(compute_one::<S, F>(), Fixed::<S, F>::compute_multiply(y, y))
.expect("hyperbolic metric: 1/y^2 exceeds the compute tier");
let scale = Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(scale));
let z = Fixed::ZERO;
FixedMatrix::<Fixed<S, F>>::from_slice(2, 2, &[
scale, z,
z, scale,
])
}
fn christoffel_closed_form(&self, p: &FixedVector<Fixed<S, F>>) -> Option<Tensor<Fixed<S, F>>> {
let gamma = self.christoffel_closed_form_compute(p)?;
Some(round_tensor(&[2, 2, 2], &gamma.data).expect("hyperbolic: Christoffel symbol exceeds storage"))
}
#[doc(hidden)]
fn christoffel_closed_form_compute(&self, p: &FixedVector<Fixed<S, F>>) -> Option<ComputeChristoffel<Fixed<S, F>>> {
let y = p[1];
if y.is_zero() { return None; }
let inv_y = Fixed::<S, F>::compute_divide(compute_one::<S, F>(), Fixed::<S, F>::upscale_to_compute(y.raw()))
.expect("hyperbolic: 1/y exceeds the compute tier");
let mut data = vec![compute_zero::<S, F>(); 8];
data[1] = Fixed::<S, F>::compute_negate(inv_y);
data[2] = Fixed::<S, F>::compute_negate(inv_y);
data[4] = inv_y;
data[7] = Fixed::<S, F>::compute_negate(inv_y);
Some(ComputeChristoffel::<Fixed<S, F>> { data })
}
fn scalar_curvature_closed_form(&self, _p: &FixedVector<Fixed<S, F>>) -> Option<Fixed<S, F>> {
Some(Fixed::from_int(-2))
}
}
pub struct GeodesicOde<'a, T> {
provider: &'a dyn MetricProvider<T>,
}
impl<'a, S: Storage, const F: u32> OdeSystem<Fixed<S, F>> for GeodesicOde<'a, Fixed<S, F>>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
fn eval(&self, _t: Fixed<S, F>, state: &FixedVector<Fixed<S, F>>) -> FixedVector<Fixed<S, F>> {
let n = self.provider.dimension();
let mut x = FixedVector::<Fixed<S, F>>::new(n);
let mut v = FixedVector::<Fixed<S, F>>::new(n);
for i in 0..n { x[i] = state[i]; v[i] = state[n + i]; }
let gamma = match christoffel_compute(self.provider, &x) {
Ok(g) => g,
Err(_) => return FixedVector::<Fixed<S, F>>::new(2 * n), };
let mut dstate = FixedVector::<Fixed<S, F>>::new(2 * n);
for k in 0..n { dstate[k] = v[k]; }
let v_c = state_to_compute(&v);
for k in 0..n {
let mut acc = Fixed::<S, F>::upscale_to_compute(Fixed::ZERO.raw());
for i in 0..n {
for j in 0..n {
let vv = Fixed::<S, F>::compute_multiply(v_c[i], v_c[j]);
acc = Fixed::<S, F>::compute_add(acc, Fixed::<S, F>::compute_multiply(gamma[(k * n + i) * n + j], vv));
}
}
dstate[n + k] = Fixed::from_raw(kernels::round_to_storage::<Fixed<S, F>>(Fixed::<S, F>::compute_negate(acc)));
}
dstate
}
}
pub fn geodesic_integrate<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
initial_point: &FixedVector<Fixed<S, F>>,
initial_velocity: &FixedVector<Fixed<S, F>>,
total_time: Fixed<S, F>,
num_steps: usize,
) -> Result<Vec<FixedVector<Fixed<S, F>>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
let n = provider.dimension();
let time_at = |k: usize| -> Compute<Fixed<S, F>> {
Fixed::<S, F>::compute_mul_div_int(Fixed::<S, F>::upscale_to_compute(total_time.raw()), k as i64, num_steps as i64)
.expect("geodesic: num_steps > 0 and the time fits")
};
let mut state = FixedVector::<Fixed<S, F>>::new(2 * n);
for i in 0..n { state[i] = initial_point[i]; state[n + i] = initial_velocity[i]; }
let mut state = state_to_compute(&state);
let sys = GeodesicOde::<Fixed<S, F>> { provider };
let mut points = Vec::with_capacity(num_steps + 1);
let mut t = time_at(0);
let extract_pos = |s: &[Compute<Fixed<S, F>>]| -> Result<FixedVector<Fixed<S, F>>, OverflowDetected> {
let mut p = FixedVector::<Fixed<S, F>>::new(n);
for i in 0..n { p[i] = Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(s[i])?); }
Ok(p)
};
points.push(initial_point.clone());
for k in 0..num_steps {
let t_next = time_at(k + 1);
state = rk4_step_compute(&sys, t, &state, Fixed::<S, F>::compute_subtract(t_next, t));
t = t_next;
points.push(extract_pos(&state)?);
}
Ok(points)
}
pub fn parallel_transport_ode<S: Storage, const F: u32>(
provider: &dyn MetricProvider<Fixed<S, F>>,
curve: &[FixedVector<Fixed<S, F>>],
initial_vector: &FixedVector<Fixed<S, F>>,
reorthog_interval: usize,
) -> Result<FixedVector<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
if curve.len() < 2 {
return Ok(initial_vector.clone());
}
let n = provider.dimension();
let mut v = state_to_compute(initial_vector);
let zero = Fixed::<S, F>::upscale_to_compute(Fixed::ZERO.raw());
for step in 0..curve.len() - 1 {
let p = &curve[step];
let p_next = &curve[step + 1];
let dx: Vec<Fixed<S, F>> = (0..n).map(|i| p_next[i] - p[i]).collect();
let gamma = christoffel_compute(provider, p)?;
let dx_c: Vec<Compute<Fixed<S, F>>> = dx.iter().map(|d| Fixed::<S, F>::upscale_to_compute(d.raw())).collect();
let mut v_new = Vec::with_capacity(n);
for k in 0..n {
let mut correction = zero;
for j in 0..n {
let mut gamma_v = zero;
for i in 0..n {
gamma_v = Fixed::<S, F>::compute_checked_add(gamma_v, kernels::compute_product::<Fixed<S, F>>(gamma[(k * n + i) * n + j], v[i])?)?;
}
correction = Fixed::<S, F>::compute_checked_add(correction, kernels::compute_product::<Fixed<S, F>>(gamma_v, dx_c[j])?)?;
}
v_new.push(checked_sub(v[k], correction)?);
}
v = v_new;
if reorthog_interval > 0 && (step + 1) % reorthog_interval == 0 {
let mut dx_norm_sq = zero;
let mut v_dot_dx = zero;
for i in 0..n {
dx_norm_sq = Fixed::<S, F>::compute_checked_add(dx_norm_sq, kernels::compute_product::<Fixed<S, F>>(dx_c[i], dx_c[i])?)?;
v_dot_dx = Fixed::<S, F>::compute_checked_add(v_dot_dx, kernels::compute_product::<Fixed<S, F>>(v[i], dx_c[i])?)?;
}
if dx_norm_sq != zero {
let coeff = Fixed::<S, F>::compute_checked_divide(v_dot_dx, dx_norm_sq)?;
for i in 0..n {
v[i] = checked_sub(v[i], kernels::compute_product::<Fixed<S, F>>(dx_c[i], coeff)?)?;
}
}
}
}
let mut out = FixedVector::<Fixed<S, F>>::new(n);
for i in 0..n { out[i] = Fixed::from_raw(Fixed::<S, F>::downscale_to_storage(v[i])?); }
Ok(out)
}
fn checked_sub<S: Storage, const F: u32>(a: Compute<Fixed<S, F>>, b: Compute<Fixed<S, F>>) -> Result<Compute<Fixed<S, F>>, OverflowDetected>
where
Fixed<S, F>: FixedNum<Storage = S>,
{
a.checked_sub_wide(b).ok_or(OverflowDetected::TierOverflow)
}