use std::iter::Sum;
use num_traits::{AsPrimitive, Float};
use rayon::prelude::*;
use rustfft::{FftNum, num_complex::Complex};
use super::{fft::FftGrid, morton::Dim};
#[doc(hidden)]
pub trait FftDim {}
impl FftDim for Dim<1> {}
impl FftDim for Dim<2> {}
const INTERPOLATION_POINTS: usize = 3;
const MIN_BOXES_PER_AXIS: usize = 10;
const EXTENT_PER_BOX: f64 = 1.0;
const MAX_BOXES_PER_AXIS: usize = 1000;
pub(crate) struct Interpolant<T: FftNum, const D: usize> {
nodes_per_axis: usize,
fft_len: usize,
plan: FftGrid<T>,
real_grid: Vec<T>,
kernel_spec: Vec<Complex<T>>,
work_spec: Vec<Complex<T>>,
potentials: Vec<T>,
box_index: Vec<usize>,
weights: Vec<T>,
}
impl<T, const D: usize> Interpolant<T, D>
where
T: Send + Sync + Float + FftNum + AsPrimitive<usize> + Sum,
{
const TERMS: usize = const { D + 2 };
pub(crate) fn new() -> Self {
Self {
nodes_per_axis: 0,
fft_len: 1,
plan: FftGrid::new(1),
real_grid: Vec::new(),
kernel_spec: Vec::new(),
work_spec: Vec::new(),
potentials: Vec::new(),
box_index: Vec::new(),
weights: Vec::new(),
}
}
pub(crate) fn repulsive_forces(
&mut self,
y: &[T],
n_samples: usize,
negative_forces: &mut [T],
z: &mut T,
) {
if n_samples == 0 {
*z = T::zero();
return;
}
let (min, max) = coordinate_bounds(y);
let extent = (max - min).max(T::min_positive_value());
let boxes_per_axis = box_count(extent);
let nodes_per_axis = boxes_per_axis * INTERPOLATION_POINTS;
let fft_len = next_smooth(2 * nodes_per_axis - 1);
self.resize(nodes_per_axis, fft_len, n_samples);
let box_width = extent / T::from(boxes_per_axis).unwrap();
let spacing = box_width / T::from(INTERPOLATION_POINTS).unwrap();
self.locate_points(y, min, box_width, boxes_per_axis);
self.build_kernel_spectrum(spacing);
self.convolve_terms(y);
self.reconstruct(y, n_samples, negative_forces, z);
}
fn resize(&mut self, nodes_per_axis: usize, fft_len: usize, n_samples: usize) {
self.nodes_per_axis = nodes_per_axis;
if fft_len != self.fft_len {
self.fft_len = fft_len;
self.plan = FftGrid::new(fft_len);
}
let grid = fft_len.pow(D as u32);
let spectrum = self.plan.spectrum_len::<D>();
let zero = Complex::new(T::zero(), T::zero());
self.real_grid.resize(grid, T::zero());
self.kernel_spec.resize(spectrum, zero);
self.work_spec.resize(spectrum, zero);
self.potentials.resize(n_samples * Self::TERMS, T::zero());
self.box_index.resize(n_samples * D, 0);
self.weights
.resize(n_samples * D * INTERPOLATION_POINTS, T::zero());
}
fn locate_points(&mut self, y: &[T], min: T, box_width: T, boxes_per_axis: usize) {
let inv_box_width = box_width.recip();
let last_box = boxes_per_axis - 1;
let basis = LagrangeBasis::<INTERPOLATION_POINTS, T>::new();
let (box_index_chunks, _) = self.box_index.as_chunks_mut::<D>();
let (y_chunks, _) = y.as_chunks::<D>();
box_index_chunks
.par_iter_mut()
.zip(self.weights.par_chunks_mut(D * INTERPOLATION_POINTS))
.zip(y_chunks.par_iter())
.for_each(|((box_row, weight_row), point)| {
for axis in 0..D {
let offset = (point[axis] - min) * inv_box_width;
let box_id = offset.floor();
let box_id = if box_id > T::zero() {
box_id.to_usize().unwrap_or(last_box).min(last_box)
} else {
0
};
box_row[axis] = box_id;
let local = offset - T::from(box_id).unwrap();
basis.weights(local, &mut weight_row[axis * INTERPOLATION_POINTS..]);
}
});
}
fn build_kernel_spectrum(&mut self, spacing: T) {
let fft_len = self.fft_len;
let one = T::one();
self.real_grid
.par_iter_mut()
.enumerate()
.for_each(|(flat, value)| {
let mut sq = T::zero();
let mut rem = flat;
for _ in 0..D {
let coord = rem % fft_len;
rem /= fft_len;
let lag = signed_lag(coord, fft_len);
let displacement = spacing * T::from(lag).unwrap();
sq = sq + displacement * displacement;
}
let cauchy = (one + sq).recip();
*value = cauchy * cauchy;
});
self.plan
.forward::<D>(&mut self.real_grid, &mut self.kernel_spec);
let inv_total = T::from(fft_len.pow(D as u32)).unwrap().recip();
self.kernel_spec
.par_iter_mut()
.for_each(|value| *value = *value * inv_total);
}
fn convolve_terms(&mut self, y: &[T]) {
for term in 0..Self::TERMS {
self.spread_charges(y, term);
self.plan
.forward::<D>(&mut self.real_grid, &mut self.work_spec);
self.multiply_by_kernel();
self.plan
.inverse::<D>(&mut self.work_spec, &mut self.real_grid);
self.gather_potentials(term);
}
}
fn spread_charges(&mut self, y: &[T], term: usize) {
let fft_len = self.fft_len;
let combinations = INTERPOLATION_POINTS.pow(D as u32);
let real_grid = &mut self.real_grid;
let (y_chunks, _) = y.as_chunks::<D>();
let (box_chunks, _) = self.box_index.as_chunks::<D>();
real_grid.iter_mut().for_each(|v| *v = T::zero());
for ((point, box_row), weight_row) in y_chunks
.iter()
.zip(box_chunks.iter())
.zip(self.weights.chunks_exact(D * INTERPOLATION_POINTS))
{
let charge = charge_value::<T, D>(point, term);
for combo in 0..combinations {
let (flat, weight) = node_of_combo::<T, D>(combo, box_row, weight_row, fft_len);
real_grid[flat] = real_grid[flat] + weight * charge;
}
}
}
fn multiply_by_kernel(&mut self) {
self.work_spec
.par_iter_mut()
.zip(self.kernel_spec.par_iter())
.for_each(|(w, &k)| *w = *w * k);
}
fn gather_potentials(&mut self, term: usize) {
let fft_len = self.fft_len;
let combinations = INTERPOLATION_POINTS.pow(D as u32);
let real_grid = &self.real_grid;
let (box_chunks, _) = self.box_index.as_chunks::<D>();
let weights = &self.weights;
self.potentials
.par_chunks_mut(Self::TERMS)
.zip(box_chunks.par_iter())
.zip(weights.par_chunks(D * INTERPOLATION_POINTS))
.for_each(|((point_terms, box_row), weight_row)| {
let mut phi = T::zero();
for combo in 0..combinations {
let (flat, weight) = node_of_combo::<T, D>(combo, box_row, weight_row, fft_len);
phi = phi + weight * real_grid[flat];
}
point_terms[term] = phi;
});
}
fn reconstruct(&self, y: &[T], n_samples: usize, negative_forces: &mut [T], z: &mut T) {
let two = T::from(2.0).unwrap();
let (negative_forces_chunks, _) = negative_forces.as_chunks_mut::<D>();
let (y_chunks, _) = y.as_chunks::<D>();
let z_sum: T = negative_forces_chunks
.par_iter_mut()
.zip(y_chunks.par_iter())
.zip(self.potentials.par_chunks(Self::TERMS))
.map(|((force_row, point), phi)| {
let density = phi[0];
let squared_norm_potential = phi[D + 1];
let mut norm_sq = T::zero();
let mut cross = T::zero();
for d in 0..D {
let first_moment = phi[1 + d];
force_row[d] = point[d] * density - first_moment;
norm_sq = norm_sq + point[d] * point[d];
cross = cross + point[d] * first_moment;
}
(T::one() + norm_sq) * density - two * cross + squared_norm_potential
})
.sum();
*z = z_sum - T::from(n_samples).unwrap();
}
}
fn coordinate_bounds<T: Float + Send + Sync>(y: &[T]) -> (T, T) {
y.par_iter()
.copied()
.fold(
|| (T::infinity(), T::neg_infinity()),
|(lo, hi), v| (lo.min(v), hi.max(v)),
)
.reduce(
|| (T::infinity(), T::neg_infinity()),
|(lo_a, hi_a), (lo_b, hi_b)| (lo_a.min(lo_b), hi_a.max(hi_b)),
)
}
fn box_count<T: Float + AsPrimitive<usize>>(extent: T) -> usize {
let scaled = (extent / T::from(EXTENT_PER_BOX).unwrap()).ceil();
let scaled = if scaled > T::zero() { scaled.as_() } else { 0 };
scaled.clamp(MIN_BOXES_PER_AXIS, MAX_BOXES_PER_AXIS)
}
fn next_smooth(n: usize) -> usize {
let mut candidate = n.max(1);
loop {
let mut value = candidate;
for prime in [2, 3, 5, 7] {
while value.is_multiple_of(prime) {
value /= prime;
}
}
if value == 1 {
return candidate;
}
candidate += 1;
}
}
#[inline]
const fn signed_lag(coord: usize, fft_len: usize) -> isize {
if coord <= fft_len / 2 {
coord as isize
} else {
coord as isize - fft_len as isize
}
}
struct LagrangeBasis<const I: usize, T> {
nodes: [T; I],
inv_denom: [T; I],
}
impl<const I: usize, T: Float> LagrangeBasis<I, T> {
fn new() -> Self {
let half = T::from(0.5).unwrap();
let p_recip = T::from(I).unwrap().recip();
let mut nodes = [T::zero(); I];
for (k, node_k) in nodes.iter_mut().enumerate() {
*node_k = (T::from(k).unwrap() + half) * p_recip;
}
let mut inv_denom = [T::one(); I];
for (k, inv_k) in inv_denom.iter_mut().enumerate() {
let mut denom = T::one();
for (m, &node_m) in nodes.iter().enumerate() {
if m != k {
denom = denom * (nodes[k] - node_m);
}
}
*inv_k = denom.recip();
}
Self { nodes, inv_denom }
}
#[inline]
fn weights(&self, local: T, out: &mut [T]) {
for (k, out_k) in out[..I].iter_mut().enumerate() {
let mut num = T::one();
for (m, &node_m) in self.nodes.iter().enumerate() {
if m != k {
num = num * (local - node_m);
}
}
*out_k = num * self.inv_denom[k];
}
}
}
#[inline]
fn charge_value<T: Float + Sum, const D: usize>(point: &[T; D], term: usize) -> T {
if term == 0 {
return T::one();
}
if (1..=D).contains(&term) {
return point[term - 1];
}
point.iter().map(|&c| c * c).sum()
}
#[inline]
fn node_of_combo<T: Float, const D: usize>(
combo: usize,
box_row: &[usize; D],
weight_row: &[T],
fft_len: usize,
) -> (usize, T) {
let p = INTERPOLATION_POINTS;
let mut rem = combo;
let mut flat = 0usize;
let mut weight = T::one();
for axis in 0..D {
let k = rem % p;
rem /= p;
let node = box_row[axis] * p + k;
flat = flat * fft_len + node;
weight = weight * weight_row[axis * p + k];
}
(flat, weight)
}
#[cfg(test)]
mod tests {
use rand::{Rng, SeedableRng, rngs::StdRng};
use super::*;
fn brute_force<const D: usize>(y: &[f64], n: usize) -> (Vec<f64>, f64) {
let mut forces = vec![0.0; n * D];
let mut z = 0.0;
for i in 0..n {
for j in 0..n {
if i == j {
continue;
}
let mut dist_sq = 0.0;
let mut delta = [0.0; D];
for d in 0..D {
delta[d] = y[i * D + d] - y[j * D + d];
dist_sq += delta[d] * delta[d];
}
let cauchy = 1.0 / (1.0 + dist_sq);
z += cauchy;
let sq = cauchy * cauchy;
for d in 0..D {
forces[i * D + d] += sq * delta[d];
}
}
}
(forces, z)
}
#[test]
fn repulsive_forces_match_brute_force_2d() {
const D: usize = 2;
let n = 800;
let mut rng = StdRng::seed_from_u64(0xBADC0DE);
let y: Vec<f64> = (0..n * D).map(|_| rng.random_range(-5.0..5.0)).collect();
let (ref_forces, ref_z) = brute_force::<D>(&y, n);
let mut interpolant = Interpolant::<f64, D>::new();
let mut forces = vec![0.0; n * D];
let mut z = 0.0;
interpolant.repulsive_forces(&y, n, &mut forces, &mut z);
assert!(
(z - ref_z).abs() / ref_z < 1e-3,
"Z relative error too large: got {z}, expected {ref_z}"
);
let mut diff_sq = 0.0;
let mut ref_sq = 0.0;
for k in 0..n * D {
diff_sq += (forces[k] - ref_forces[k]).powi(2);
ref_sq += ref_forces[k].powi(2);
}
let rel = (diff_sq / ref_sq).sqrt();
assert!(rel < 3e-2, "force relative L2 error too large: {rel}");
}
#[test]
fn repulsive_forces_match_brute_force_1d() {
const D: usize = 1;
let n = 500;
let mut rng = StdRng::seed_from_u64(0xFEED_F00D);
let y: Vec<f64> = (0..n * D).map(|_| rng.random_range(-4.0..4.0)).collect();
let (ref_forces, ref_z) = brute_force::<D>(&y, n);
let mut interpolant = Interpolant::<f64, D>::new();
let mut forces = vec![0.0; n * D];
let mut z = 0.0;
interpolant.repulsive_forces(&y, n, &mut forces, &mut z);
assert!((z - ref_z).abs() / ref_z < 1e-3, "Z error: {z} vs {ref_z}");
let mut diff_sq = 0.0;
let mut ref_sq = 0.0;
for k in 0..n * D {
diff_sq += (forces[k] - ref_forces[k]).powi(2);
ref_sq += ref_forces[k].powi(2);
}
assert!((diff_sq / ref_sq).sqrt() < 3e-2);
}
#[test]
fn lagrange_weights_partition_of_unity() {
let basis = LagrangeBasis::<INTERPOLATION_POINTS, f64>::new();
for step in 0..=10 {
let local = step as f64 / 10.0;
let mut w = [0.0; INTERPOLATION_POINTS];
basis.weights(local, &mut w);
let sum: f64 = w.iter().sum();
assert!((sum - 1.0).abs() < 1e-12, "weights sum to {sum} at {local}");
}
for j in 0..INTERPOLATION_POINTS {
let node_j = (j as f64 + 0.5) / INTERPOLATION_POINTS as f64;
let mut w = [0.0; INTERPOLATION_POINTS];
basis.weights(node_j, &mut w);
for (k, &w_k) in w.iter().enumerate() {
let expected = if k == j { 1.0 } else { 0.0 };
assert!((w_k - expected).abs() < 1e-12, "w[{k}] = {w_k} at node {j}");
}
}
}
#[test]
fn empty_input_is_a_no_op() {
let mut interpolant = Interpolant::<f64, 2>::new();
let mut forces: Vec<f64> = Vec::new();
let mut z = 1.0;
interpolant.repulsive_forces(&[], 0, &mut forces, &mut z);
assert_eq!(z, 0.0);
}
}