use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Sub, SubAssign};
use na::Const;
use na::Scalar;
use nalgebra as na;
pub fn gradient<T, const N: usize>(
x: &na::Matrix<T, Const<N>, Const<1>, na::ArrayStorage<T, N, 1>>,
f: &dyn Fn(&na::Matrix<T, Const<N>, Const<1>, na::ArrayStorage<T, N, 1>>) -> T,
epsilon: T,
) -> na::Matrix<T, Const<N>, Const<1>, na::ArrayStorage<T, N, 1>>
where
T: Scalar
+ Add<Output = T>
+ Sub<Output = T>
+ Div<Output = T>
+ Mul<Output = T>
+ Copy,
na::ArrayStorage<T, N, 1>: Default,
{
let mut grad =
na::Matrix::<T, Const<N>, Const<1>, na::ArrayStorage<T, N, 1>>::default();
for i in 0..N
{
let mut nx = x.clone();
nx[i] = nx[i] + epsilon;
let mut px = x.clone();
px[i] = px[i] - epsilon;
grad[i] = (f(&nx) - f(&px)) / (epsilon + epsilon);
}
grad
}
pub fn divergence<T, const N: usize>(
x: &na::Matrix<T, Const<N>, Const<1>, na::ArrayStorage<T, N, 1>>,
f: &dyn Fn(
&na::Matrix<T, Const<N>, Const<1>, na::ArrayStorage<T, N, 1>>,
) -> na::Matrix<T, Const<N>, Const<1>, na::ArrayStorage<T, N, 1>>,
epsilon: T,
) -> T
where
T: Scalar
+ Add<Output = T>
+ Sub<Output = T>
+ Div<Output = T>
+ Mul<Output = T>
+ Copy
+ Default,
na::ArrayStorage<T, N, 1>: Default,
{
let mut div = T::default();
for i in 0..N
{
let mut nx = x.clone();
nx[i] = nx[i] + epsilon;
let mut px = x.clone();
px[i] = px[i] - epsilon;
div = div + (f(&nx)[i] - f(&px)[i]) / (epsilon + epsilon);
}
div
}
type Vec3<T> = na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>;
pub fn mean_curvature<T>(
x: &na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>,
f: &dyn Fn(&na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>) -> T,
epsilon: T,
) -> T
where
T: Scalar
+ Add<Output = T>
+ Sub<Output = T>
+ Div<Output = T>
+ Mul<Output = T>
+ Copy
+ Default
+ AddAssign
+ SubAssign
+ num_traits::Float,
na::ArrayStorage<T, 3, 1>: Default,
{
let z = T::from(0.).unwrap();
let t1 = f(&(x + Vec3::new(epsilon, z, z)));
let t2 = f(&(x - Vec3::new(epsilon, z, z)));
let t3 = f(&(x + Vec3::new(z, epsilon, z)));
let t4 = f(&(x - Vec3::new(z, epsilon, z)));
let t5 = f(&(x + Vec3::new(z, z, epsilon)));
let t6 = f(&(x - Vec3::new(z, z, epsilon)));
(T::from(0.5).unwrap() / (epsilon * epsilon))
* (t1 + t2 + t3 + t4 + t5 + t6 - T::from(6.).unwrap() * f(x))
}
pub fn gaussian_curvature<T>(
x: &na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>,
f: &dyn Fn(&na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>) -> T,
epsilon: T,
) -> T
where
T: Scalar
+ Add<Output = T>
+ Sub<Output = T>
+ Div<Output = T>
+ Mul<Output = T>
+ Copy
+ Default
+ AddAssign
+ SubAssign
+ num_traits::Float
+ MulAssign,
na::ArrayStorage<T, 3, 1>: Default,
{
let g = gradient(x, f, epsilon);
let h = adjoint_hessian(x, f, epsilon);
let k = g.transpose() * h * g;
let n = g.dot(&g);
let curv = k[(0, 0)] / (n * n);
curv
}
pub fn hessian<T>(
x: &na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>,
f: &dyn Fn(&na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>) -> T,
epsilon: T,
) -> na::Matrix3<T>
where
T: Scalar
+ Add<Output = T>
+ Sub<Output = T>
+ Div<Output = T>
+ Mul<Output = T>
+ Copy
+ Default
+ AddAssign
+ SubAssign
+ num_traits::Float,
na::ArrayStorage<T, 3, 1>: Default,
{
let mut hessian = na::Matrix3::zeros();
for i in 0..3
{
for j in 0..3
{
let mut acc = T::from(0.).unwrap();
for dx in (-1..2).step_by(2)
{
let ox = T::from(dx).unwrap() * epsilon;
for dy in (-1..2).step_by(2)
{
let oy = T::from(dy).unwrap() * epsilon;
let mut sample = x.clone();
sample[i] += ox;
sample[j] += oy;
acc += f(&sample);
}
}
hessian[(i, j)] = acc / (T::from(4.).unwrap() * epsilon * epsilon);
}
}
hessian
}
pub fn adjoint_hessian<T>(
x: &na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>,
f: &dyn Fn(&na::Matrix<T, Const<3>, Const<1>, na::ArrayStorage<T, 3, 1>>) -> T,
epsilon: T,
) -> na::Matrix3<T>
where
T: Scalar
+ Add<Output = T>
+ Sub<Output = T>
+ Div<Output = T>
+ Mul<Output = T>
+ Copy
+ Default
+ AddAssign
+ SubAssign
+ num_traits::Float,
na::ArrayStorage<T, 3, 1>: Default,
{
let h = hessian(x, f, epsilon);
let mut adjoint = na::Matrix3::zeros();
adjoint[(0, 0)] = h[(1, 1)] * h[(2, 2)] - h[(1, 2)] * h[(2, 1)];
adjoint[(0, 1)] = h[(1, 2)] * h[(2, 0)] - h[(1, 0)] * h[(2, 2)];
adjoint[(0, 2)] = h[(1, 0)] * h[(2, 1)] - h[(1, 1)] * h[(2, 0)];
adjoint[(1, 0)] = h[(0, 2)] * h[(2, 1)] - h[(0, 1)] * h[(2, 2)];
adjoint[(1, 1)] = h[(0, 0)] * h[(2, 2)] - h[(0, 2)] * h[(2, 0)];
adjoint[(1, 2)] = h[(0, 1)] * h[(2, 0)] - h[(0, 0)] * h[(2, 1)];
adjoint[(2, 0)] = h[(0, 1)] * h[(1, 2)] - h[(0, 2)] * h[(1, 1)];
adjoint[(2, 1)] = h[(1, 0)] * h[(0, 2)] - h[(0, 0)] * h[(1, 2)];
adjoint[(2, 2)] = h[(0, 0)] * h[(1, 1)] - h[(0, 1)] * h[(1, 0)];
adjoint
}