use crate::brute::{canonicalise, distance_squared, finite};
use crate::{CellGridOptions, NeighborPair, PeriodicBox, SpatialError};
const NEIGHBOUR_OFFSETS: [[isize; 3]; 27] = super::NEIGHBOUR_OFFSETS;
#[derive(Debug)]
pub(super) struct PeriodicGrid {
dims: [usize; 3],
offsets: Vec<usize>,
members: Vec<u32>,
}
impl PeriodicGrid {
pub(super) fn retained_bytes(&self) -> usize {
self.offsets.capacity() * size_of::<usize>() + self.members.capacity() * size_of::<u32>()
}
pub(super) fn build(
positions: &[[f32; 3]],
targets: &[u32],
cutoff: f32,
periodic: &PeriodicBox,
options: CellGridOptions,
) -> Result<Self, SpatialError> {
let dims = periodic_dimensions(periodic, cutoff, options)?;
let cell_count = cell_count(dims)?;
let mut counts = vec![0usize; cell_count];
for &target in targets {
let index = usize::try_from(target).map_err(|_| SpatialError::NumericRangeExceeded)?;
let Some(position) = positions.get(index).copied() else {
continue;
};
if finite(position) {
counts[linear(dims, cell(periodic, position, dims)?)?] += 1;
}
}
let offsets = prefix_offsets(&counts)?;
counts.copy_from_slice(&offsets[..cell_count]);
let finite_count = match offsets.last() {
Some(count) => *count,
None => 0,
};
let mut members = vec![0u32; finite_count];
for &target in targets {
let index = usize::try_from(target).map_err(|_| SpatialError::NumericRangeExceeded)?;
let Some(position) = positions.get(index).copied() else {
continue;
};
if !finite(position) {
continue;
}
let index = linear(dims, cell(periodic, position, dims)?)?;
let Some(cursor) = counts.get_mut(index) else {
continue;
};
if let Some(slot) = members.get_mut(*cursor) {
*slot = target;
*cursor += 1;
}
}
Ok(Self {
dims,
offsets,
members,
})
}
pub(super) fn for_each_pair<F>(
&self,
positions: &[[f32; 3]],
query: &[u32],
cutoff_squared: f32,
periodic: &PeriodicBox,
emit: &mut F,
) -> Result<(), SpatialError>
where
F: FnMut(u32, u32, f32),
{
for_each_query_position(positions, query, |atom, position| {
self.append_pairs(positions, atom, position, cutoff_squared, periodic, emit)
})
}
pub(super) fn pairs_same_selection(
&self,
positions: &[[f32; 3]],
query: &[u32],
cutoff_squared: f32,
periodic: &PeriodicBox,
sort_result: bool,
) -> Result<Vec<NeighborPair>, SpatialError> {
let mut found = Vec::new();
self.for_each_pairs_same_selection(
positions,
query,
cutoff_squared,
periodic,
&mut |pair| found.push(pair),
)?;
if sort_result {
canonicalise(&mut found);
}
Ok(found)
}
pub(super) fn for_each_pairs_same_selection<F>(
&self,
positions: &[[f32; 3]],
query: &[u32],
cutoff_squared: f32,
periodic: &PeriodicBox,
emit: &mut F,
) -> Result<(), SpatialError>
where
F: FnMut(NeighborPair),
{
for_each_query_position(positions, query, |atom, position| {
self.append_pairs_same_selection(
positions,
atom,
position,
cutoff_squared,
periodic,
emit,
)
})
}
fn append_pairs<F>(
&self,
positions: &[[f32; 3]],
atom: u32,
position: [f32; 3],
cutoff_squared: f32,
periodic: &PeriodicBox,
emit: &mut F,
) -> Result<(), SpatialError>
where
F: FnMut(u32, u32, f32),
{
self.for_each_neighbor(
positions,
position,
cutoff_squared,
periodic,
&mut |target, squared| {
if atom != target {
emit(atom, target, squared);
}
},
)
}
pub(super) fn for_each_neighbor(
&self,
positions: &[[f32; 3]],
position: [f32; 3],
cutoff_squared: f32,
periodic: &PeriodicBox,
emit: &mut impl FnMut(u32, f32),
) -> Result<(), SpatialError> {
let centre = cell(periodic, position, self.dims)?;
let mut visited = [usize::MAX; NEIGHBOUR_OFFSETS.len()];
let mut visited_count = 0usize;
for delta in NEIGHBOUR_OFFSETS {
let candidate = periodic_neighbour(self.dims, centre, delta)?;
if visited[..visited_count].contains(&candidate) {
continue;
}
visited[visited_count] = candidate;
visited_count += 1;
let members = self
.members(candidate)
.ok_or(SpatialError::NumericRangeExceeded)?;
for &target in members {
let index =
usize::try_from(target).map_err(|_| SpatialError::NumericRangeExceeded)?;
let Some(target_position) = positions.get(index).copied() else {
continue;
};
let squared = distance_squared(position, target_position, Some(periodic));
if squared <= cutoff_squared {
emit(target, squared);
}
}
}
Ok(())
}
fn append_pairs_same_selection<F>(
&self,
positions: &[[f32; 3]],
atom: u32,
position: [f32; 3],
cutoff_squared: f32,
periodic: &PeriodicBox,
emit: &mut F,
) -> Result<(), SpatialError>
where
F: FnMut(NeighborPair),
{
let centre = cell(periodic, position, self.dims)?;
let mut visited = [usize::MAX; NEIGHBOUR_OFFSETS.len()];
let mut visited_count = 0usize;
for delta in NEIGHBOUR_OFFSETS {
let candidate = periodic_neighbour(self.dims, centre, delta)?;
if visited[..visited_count].contains(&candidate) {
continue;
}
visited[visited_count] = candidate;
visited_count += 1;
let members = self
.members(candidate)
.ok_or(SpatialError::NumericRangeExceeded)?;
for &target in members {
if atom >= target {
continue;
}
let index =
usize::try_from(target).map_err(|_| SpatialError::NumericRangeExceeded)?;
let Some(target_position) = positions.get(index).copied() else {
continue;
};
let squared = distance_squared(position, target_position, Some(periodic));
if squared <= cutoff_squared {
emit(NeighborPair::new(atom, target, squared));
}
}
}
Ok(())
}
fn members(&self, cell: usize) -> Option<&[u32]> {
let &start = self.offsets.get(cell)?;
let next = cell.checked_add(1)?;
let &end = self.offsets.get(next)?;
self.members.get(start..end)
}
}
pub(super) fn periodic_dimensions(
periodic: &PeriodicBox,
cutoff: f32,
options: CellGridOptions,
) -> Result<[usize; 3], SpatialError> {
if cutoff == 0.0 {
return Ok([1; 3]);
}
let mut reach = periodic.fractional_cutoff_bounds(cutoff);
let mut dims = dimensions(reach, options.maximum_cell_count)?;
while cell_count(dims)? > options.maximum_cell_count {
reach = reach.map(|value| value * f64::from(options.edge_growth_factor));
dims = dimensions(reach, options.maximum_cell_count)?;
}
Ok(dims)
}
fn dimensions(reach: [f64; 3], maximum: usize) -> Result<[usize; 3], SpatialError> {
Ok([
dimension_for_reach(reach[0], maximum)?,
dimension_for_reach(reach[1], maximum)?,
dimension_for_reach(reach[2], maximum)?,
])
}
fn dimension_for_reach(reach: f64, maximum: usize) -> Result<usize, SpatialError> {
if reach <= 0.0 {
Ok(1)
} else {
let value =
crate::numeric::floor_usize(reach.recip()).ok_or(SpatialError::NumericRangeExceeded)?;
Ok(value.min(maximum).max(1))
}
}
fn cell(
periodic: &PeriodicBox,
position: [f32; 3],
dims: [usize; 3],
) -> Result<[usize; 3], SpatialError> {
let fractional = periodic
.fractional(position)
.map(|value| value.rem_euclid(1.0));
let coordinate = |axis: usize| -> Result<usize, SpatialError> {
let dimension =
crate::numeric::usize_f64(dims[axis]).ok_or(SpatialError::NumericRangeExceeded)?;
let value = crate::numeric::floor_usize(fractional[axis] * dimension)
.ok_or(SpatialError::NumericRangeExceeded)?;
let maximum = dims[axis].checked_sub(1).ok_or(SpatialError::InvalidCell)?;
Ok(value.min(maximum))
};
Ok([coordinate(0)?, coordinate(1)?, coordinate(2)?])
}
fn periodic_neighbour(
dims: [usize; 3],
centre: [usize; 3],
delta: [isize; 3],
) -> Result<usize, SpatialError> {
let coordinate = |axis| {
let dimension =
isize::try_from(dims[axis]).map_err(|_| SpatialError::NumericRangeExceeded)?;
let origin =
isize::try_from(centre[axis]).map_err(|_| SpatialError::NumericRangeExceeded)?;
let shifted = origin
.checked_add(delta[axis])
.ok_or(SpatialError::NumericRangeExceeded)?;
usize::try_from(shifted.rem_euclid(dimension))
.map_err(|_| SpatialError::NumericRangeExceeded)
};
linear(dims, [coordinate(0)?, coordinate(1)?, coordinate(2)?])
}
#[inline]
fn linear(dims: [usize; 3], cell: [usize; 3]) -> Result<usize, SpatialError> {
cell[0]
.checked_mul(dims[1])
.and_then(|value| value.checked_add(cell[1]))
.and_then(|value| value.checked_mul(dims[2]))
.and_then(|value| value.checked_add(cell[2]))
.ok_or(SpatialError::NumericRangeExceeded)
}
fn prefix_offsets(counts: &[usize]) -> Result<Vec<usize>, SpatialError> {
let capacity = counts
.len()
.checked_add(1)
.ok_or(SpatialError::NumericRangeExceeded)?;
let mut offsets = Vec::with_capacity(capacity);
let mut total = 0usize;
offsets.push(total);
for &count in counts {
total = total
.checked_add(count)
.ok_or(SpatialError::NumericRangeExceeded)?;
offsets.push(total);
}
Ok(offsets)
}
pub(super) fn cell_count(dims: [usize; 3]) -> Result<usize, SpatialError> {
dims[0]
.checked_mul(dims[1])
.and_then(|value| value.checked_mul(dims[2]))
.ok_or(SpatialError::NumericRangeExceeded)
}
fn for_each_query_position(
positions: &[[f32; 3]],
query: &[u32],
mut visit: impl FnMut(u32, [f32; 3]) -> Result<(), SpatialError>,
) -> Result<(), SpatialError> {
for &atom in query {
let index = usize::try_from(atom).map_err(|_| SpatialError::NumericRangeExceeded)?;
let Some(position) = positions.get(index).copied() else {
continue;
};
if finite(position) {
visit(atom, position)?;
}
}
Ok(())
}