#![warn(missing_docs)]
pub mod kernel;
pub use kernel::*;
use na::storage::Storage;
use na::{
DMatrix, DVector, Matrix3, Matrix3x4, Matrix4, Point3, RealField, Vector, Vector3, Vector4, U1,
U3,
};
use num_traits::{Float, Zero};
pub trait Real: Float + RealField + std::fmt::Debug {}
impl<T> Real for T where T: Float + RealField + std::fmt::Debug {}
pub type Pow3Hrbf<T> = Hrbf<T, kernel::Pow3<T>>;
pub type Pow5Hrbf<T> = Hrbf<T, kernel::Pow5<T>>;
pub type GaussHrbf<T> = Hrbf<T, kernel::Gauss<T>>;
pub type Csrbf31Hrbf<T> = Hrbf<T, kernel::Csrbf31<T>>;
pub type Csrbf42Hrbf<T> = Hrbf<T, kernel::Csrbf42<T>>;
pub type Pow3HrbfBuilder<T> = HrbfBuilder<T, kernel::Pow3<T>>;
pub type Pow5HrbfBuilder<T> = HrbfBuilder<T, kernel::Pow5<T>>;
pub type GaussHrbfBuilder<T> = HrbfBuilder<T, kernel::Gauss<T>>;
pub type Csrbf31HrbfBuilder<T> = HrbfBuilder<T, kernel::Csrbf31<T>>;
pub type Csrbf42HrbfBuilder<T> = HrbfBuilder<T, kernel::Csrbf42<T>>;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Error {
NumPointsMismatch,
NumOffsetsMismatch,
NumNormalsMismatch,
LinearSolveFailure,
}
impl std::fmt::Display for Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Error::NumPointsMismatch => "Number of points does not match the number of sites",
Error::NumOffsetsMismatch => "Number of offsets does not match the number of sites",
Error::NumNormalsMismatch => "Number of normals does not match the number of sites",
Error::LinearSolveFailure => "Linear solve failed when building the HRBF potential",
}
.fmt(f)
}
}
impl std::error::Error for Error {}
pub type Result<T> = std::result::Result<T, Error>;
#[derive(Clone, Debug)]
pub enum KernelType<K> {
Variable(Vec<K>),
Constant(K),
}
impl<K> ::std::ops::Index<usize> for KernelType<K> {
type Output = K;
fn index(&self, index: usize) -> &K {
match *self {
KernelType::Variable(ref ks) => &ks[index],
KernelType::Constant(ref k) => k,
}
}
}
#[derive(Clone, Debug)]
pub struct HrbfBuilder<T, K>
where
T: Real,
K: Kernel<T>,
{
sites: Vec<Point3<T>>,
points: Vec<Point3<T>>,
normals: Vec<Vector3<T>>,
offsets: Vec<T>,
kernel: KernelType<K>,
}
impl<T, K> HrbfBuilder<T, K>
where
T: Real,
K: Kernel<T> + Clone + Default,
{
pub fn new(sites: Vec<Point3<T>>) -> HrbfBuilder<T, K> {
HrbfBuilder {
sites,
points: Vec::new(),
normals: Vec::new(),
offsets: Vec::new(),
kernel: KernelType::Constant(K::default()),
}
}
pub fn offsets(&mut self, offsets: Vec<T>) -> &mut Self {
self.offsets = offsets;
self
}
pub fn points(&mut self, points: Vec<Point3<T>>) -> &mut Self {
self.points = points;
self
}
pub fn normals(&mut self, normals: Vec<Vector3<T>>) -> &mut Self {
self.normals = normals;
self
}
pub(crate) fn fit_block(
sites: &[Point3<T>],
kernel: &KernelType<K>,
p: Point3<T>,
j: usize,
) -> Matrix4<T> {
let x = p - sites[j];
let l = x.norm();
let w = kernel[j].f(l);
let g = Hrbf::grad_phi(kernel, x, l, j);
let h = Hrbf::hess_phi(kernel, x, l, j);
Matrix4::new(
w,
g[0],
g[1],
g[2],
g[0],
h[(0, 0)],
h[(0, 1)],
h[(0, 2)],
g[1],
h[(1, 0)],
h[(1, 1)],
h[(1, 2)],
g[2],
h[(2, 0)],
h[(2, 1)],
h[(2, 2)],
)
}
#[allow(non_snake_case)]
pub(crate) fn fit_system(
sites: &[Point3<T>],
points: &[Point3<T>],
offsets: &[T],
normals: &[Vector3<T>],
kernel: &KernelType<K>,
) -> (DMatrix<T>, DVector<T>) {
let num_sites = sites.len();
debug_assert_eq!(points.len(), num_sites);
debug_assert_eq!(normals.len(), num_sites);
debug_assert!(offsets.is_empty() || offsets.len() == num_sites);
let rows = 4 * points.len();
let cols = 4 * num_sites;
let mut A = DMatrix::<T>::zeros(rows, cols);
let mut b = DVector::<T>::zeros(rows);
for (i, p) in points.iter().enumerate() {
b[4 * i] = if offsets.is_empty() {
T::zero()
} else {
offsets[i]
};
b.fixed_rows_mut::<3>(4 * i + 1).copy_from(&normals[i]);
for j in 0..num_sites {
A.fixed_view_mut::<4, 4>(4 * i, 4 * j)
.copy_from(&Self::fit_block(sites, &kernel, *p, j));
}
}
(A, b)
}
pub fn build_system(&self) -> Result<(DMatrix<T>, DVector<T>)> {
let HrbfBuilder {
sites,
points,
normals,
offsets,
kernel,
} = self;
let num_sites = sites.len();
let points = if points.is_empty() {
sites.as_slice()
} else {
points.as_slice()
};
if points.len() != num_sites {
return Err(Error::NumPointsMismatch);
}
if !offsets.is_empty() || offsets.len() != num_sites {
return Err(Error::NumOffsetsMismatch);
}
if normals.len() != num_sites {
return Err(Error::NumNormalsMismatch);
}
Ok(Self::fit_system(
&sites, &points, &offsets, &normals, &kernel,
))
}
pub fn build(&self) -> Result<Hrbf<T, K>> {
let HrbfBuilder {
sites,
points,
normals,
offsets,
kernel,
} = self;
let mut hrbf = Hrbf {
sites: sites.clone(),
betas: Vec::new(),
kernel: kernel.clone(),
};
let result = if offsets.is_empty() {
if points.is_empty() {
Hrbf::fit(&mut hrbf, &normals)
} else {
Hrbf::fit_to_points(&mut hrbf, &points, &normals)
}
} else {
if points.is_empty() {
Hrbf::offset_fit(&mut hrbf, &offsets, &normals)
} else {
Hrbf::offset_fit_to_points(&mut hrbf, &points, &offsets, &normals)
}
};
match result {
Ok(_) => Ok(hrbf),
Err(e) => Err(e),
}
}
}
impl<T, K> HrbfBuilder<T, K>
where
T: Real,
K: Kernel<T> + LocalKernel<T>,
{
pub fn radius(mut self, radius: T) -> Self {
self.kernel = KernelType::Constant(K::new(radius));
self
}
pub fn radii(mut self, radii: Vec<T>) -> Self {
self.kernel = KernelType::Variable(radii.into_iter().map(|r| K::new(r)).collect());
self
}
}
#[derive(Clone, Debug)]
pub struct Hrbf<T, K>
where
T: Real,
K: Kernel<T>,
{
sites: Vec<Point3<T>>,
betas: Vec<Vector4<T>>,
kernel: KernelType<K>,
}
impl<T, K> Hrbf<T, K>
where
T: Real,
K: Kernel<T> + Clone + Default,
{
pub fn sites(&self) -> &[Point3<T>] {
&self.sites
}
pub fn betas(&self) -> &[Vector4<T>] {
&self.betas
}
#[allow(non_snake_case)]
pub fn fit(&mut self, normals: &[Vector3<T>]) -> Result<&mut Self> {
self.fit_impl(None, None, normals)
}
#[allow(non_snake_case)]
pub fn fit_to_points(
&mut self,
points: &[Point3<T>],
normals: &[Vector3<T>],
) -> Result<&mut Self> {
self.fit_impl(points.into(), None, normals)
}
#[allow(non_snake_case)]
pub fn offset_fit(&mut self, offsets: &[T], normals: &[Vector3<T>]) -> Result<&mut Self> {
self.fit_impl(None, offsets.into(), normals)
}
#[allow(non_snake_case)]
pub fn offset_fit_to_points(
&mut self,
points: &[Point3<T>],
offsets: &[T],
normals: &[Vector3<T>],
) -> Result<&mut Self> {
self.fit_impl(points.into(), offsets.into(), normals)
}
#[allow(non_snake_case)]
fn fit_impl(
&mut self,
points: Option<&[Point3<T>]>,
offsets: Option<&[T]>,
normals: &[Vector3<T>],
) -> Result<&mut Self> {
let num_sites = self.sites.len();
let points = points.unwrap_or_else(|| self.sites.as_slice());
if points.len() != num_sites {
return Err(Error::NumPointsMismatch);
}
if normals.len() != num_sites {
return Err(Error::NumNormalsMismatch);
}
let mut potential = Vec::new();
let offsets = offsets.unwrap_or_else(|| {
potential.resize(num_sites, T::zero());
potential.as_slice()
});
if offsets.len() != num_sites {
return Err(Error::NumOffsetsMismatch);
}
let (A, b) = HrbfBuilder::fit_system(
self.sites.as_slice(),
points,
offsets,
normals,
&self.kernel,
);
self.betas.clear();
if let Some(x) = A.lu().solve(&b) {
assert_eq!(x.len(), 4 * num_sites);
self.betas.resize(num_sites, Vector4::zero());
for j in 0..num_sites {
self.betas[j].copy_from(&x.fixed_rows::<4>(4 * j));
}
Ok(self)
} else {
Err(Error::LinearSolveFailure)
}
}
fn grad_phi(kernel: &KernelType<K>, x: Vector3<T>, l: T, j: usize) -> Vector3<T> {
x * kernel[j].df_l(l)
}
fn hess_phi(kernel: &KernelType<K>, x: Vector3<T>, l: T, j: usize) -> Matrix3<T> {
let df_l = kernel[j].df_l(l);
let mut hess = Matrix3::identity();
if l <= T::zero() {
debug_assert!({
let g = kernel[j].ddf(l) - df_l;
g * g < T::from(1e-12).unwrap()
});
return hess * df_l;
}
let ddf = kernel[j].ddf(l);
let x_hat = x / l;
hess.ger(ddf - df_l, &x_hat, &x_hat, df_l);
hess
}
fn third_deriv_prod_phi<S>(
&self,
x: Vector3<T>,
l: T,
b: &Vector<T, U3, S>,
j: usize,
) -> Matrix3<T>
where
S: Storage<T, U3, U1>,
{
if l <= T::zero() {
debug_assert!({
let g = self.kernel[j].g(l); let dddf = self.kernel[j].dddf(l);
dddf == T::zero() && g == T::zero()
});
return Matrix3::zero();
}
let g = self.kernel[j].g(l); let dddf = self.kernel[j].dddf(l);
let x_hat = x / l;
let x_dot_b = b.dot(&x_hat);
let mut mtx = Matrix3::identity();
let _3 = T::from(3).unwrap();
let _1 = T::one();
mtx.ger(_1, b, &x_hat, x_dot_b);
mtx.ger(_1, &x_hat, b, _1);
mtx.ger((dddf - _3 * g) * x_dot_b, &x_hat, &x_hat, g);
mtx
}
#[inline]
fn fourth_deriv_prod_phi<S>(
&self,
x: Vector3<T>,
l: T,
b: &Vector<T, U3, S>,
c: &Vector<T, U3, S>,
j: usize,
) -> Matrix3<T>
where
S: Storage<T, U3, U1>,
{
let g_l = self.kernel[j].g_l(l);
let bc_tr = Matrix3::new(
b[0] * c[0],
b[1] * c[0],
b[2] * c[0],
b[0] * c[1],
b[1] * c[1],
b[2] * c[1],
b[0] * c[2],
b[1] * c[2],
b[2] * c[2],
);
let c_dot_b = bc_tr.trace();
let bc_tr_plus_cb_tr = bc_tr + bc_tr.transpose();
let mut res = Matrix3::identity();
if l <= T::zero() {
debug_assert!({
let h3 = self.kernel[j].h(l, T::from(3).unwrap());
let h52 = self.kernel[j].h(l, T::from(5.0 / 2.0).unwrap());
let ddddf = self.kernel[j].ddddf(l);
let a = ddddf - T::from(6.0).unwrap() * h52;
h3 == T::zero() && a * a < T::from(1e-12).unwrap()
});
return res * (g_l * c_dot_b) + (bc_tr_plus_cb_tr) * g_l;
}
let h3 = self.kernel[j].h(l, T::from(3).unwrap());
let h52 = self.kernel[j].h(l, T::from(5.0 / 2.0).unwrap());
let ddddf = self.kernel[j].ddddf(l);
let a = ddddf - T::from(6.0).unwrap() * h52;
let x_hat = x / l;
let x_dot_b = x_hat.dot(b);
let x_dot_c = x_hat.dot(c);
let cb_sum = c * x_dot_b + b * x_dot_c;
let _1 = T::one();
res.ger(
a * x_dot_b * x_dot_c + h3 * c_dot_b,
&x_hat,
&x_hat,
h3 * x_dot_c * x_dot_b + g_l * c_dot_b,
);
res.ger(h3, &cb_sum, &x_hat, _1);
res.ger(h3, &x_hat, &cb_sum, _1);
res + (bc_tr_plus_cb_tr) * g_l
}
pub fn eval(&self, p: Point3<T>) -> T {
self.betas
.iter()
.enumerate()
.fold(T::zero(), |sum, (j, b)| sum + self.eval_block(p, j).dot(b))
}
fn eval_block(&self, p: Point3<T>, j: usize) -> Vector4<T> {
let x = p - self.sites[j];
let l = x.norm();
let w = self.kernel[j].f(l);
let g = Self::grad_phi(&self.kernel, x, l, j);
Vector4::new(w, g[0], g[1], g[2])
}
pub fn grad(&self, p: Point3<T>) -> Vector3<T> {
self.betas
.iter()
.enumerate()
.fold(Vector3::zero(), |sum, (j, b)| {
sum + self.grad_block(p, j) * b
})
}
fn grad_block(&self, p: Point3<T>, j: usize) -> Matrix3x4<T> {
let x = p - self.sites[j];
let l = x.norm();
let h = Self::hess_phi(&self.kernel, x, l, j);
let mut grad = Matrix3x4::zero();
grad.column_mut(0)
.copy_from(&Self::grad_phi(&self.kernel, x, l, j));
grad.fixed_columns_mut::<3>(1).copy_from(&h);
grad
}
pub fn hess(&self, p: Point3<T>) -> Matrix3<T> {
self.betas
.iter()
.enumerate()
.fold(Matrix3::zero(), |sum, (j, b)| {
sum + self.hess_block_prod(p, b, j)
})
}
#[inline]
fn hess_block_prod(&self, p: Point3<T>, b: &Vector4<T>, j: usize) -> Matrix3<T> {
let x = p - self.sites[j];
let l = x.norm();
let b3 = b.fixed_rows::<3>(1);
let h = Self::hess_phi(&self.kernel, x, l, j);
h * b[0] + self.third_deriv_prod_phi(x, l, &b3, j)
}
pub fn fit_block(&self, p: Point3<T>, j: usize) -> Matrix4<T> {
HrbfBuilder::fit_block(self.sites.as_slice(), &self.kernel, p, j)
}
pub fn grad_fit_block_prod(&self, p: Point3<T>, b: Vector4<T>, j: usize) -> Matrix3x4<T> {
let x = p - self.sites[j];
let l = x.norm();
let b3 = b.fixed_rows::<3>(1);
let g = Self::grad_phi(&self.kernel, x, l, j);
let h = Self::hess_phi(&self.kernel, x, l, j);
let third = h * b[0] + self.third_deriv_prod_phi(x, l, &b3, j);
let mut grad = Matrix3x4::zero();
grad.column_mut(0).copy_from(&(g * b[0] + h * b3));
grad.fixed_columns_mut::<3>(1).copy_from(&third);
grad
}
fn hess_fit_prod_block(&self, p: Point3<T>, c: Vector4<T>, j: usize) -> Matrix3<T> {
let x = p - self.sites[j];
let l = x.norm();
let c3 = c.fixed_rows::<3>(1);
let a = self.betas[j][0];
let b = self.betas[j].fixed_rows::<3>(1);
Self::hess_phi(&self.kernel, x, l, j) * c[0] * a
+ self.third_deriv_prod_phi(x, l, &b, j) * c[0]
+ self.third_deriv_prod_phi(x, l, &c3, j) * a
+ self.fourth_deriv_prod_phi(x, l, &b, &c3, j)
}
pub fn hess_fit_prod(&self, p: Point3<T>, c: Vector4<T>) -> Matrix3<T> {
(0..self.sites.len()).fold(Matrix3::zero(), |sum, j| {
sum + self.hess_fit_prod_block(p, c, j)
})
}
}