forceatlas2 0.8.0

use crate::{layout::*, node::*, trees::Body, util::*};

use rayon::prelude::*;

pub(crate) struct NodeBodyN<T, C, const N: usize> {
	pos: VecN<T, N>,
	mass: T,
	custom: C,
}

type NodeBody2<T, C> = NodeBodyN<T, C, 2>;
type NodeBody3<T, C> = NodeBodyN<T, C, 3>;

impl<T: Coord, C: Copy, const N: usize> Body<T, C, N> for NodeBodyN<T, C, N> {
	fn mass(&self) -> T {
		self.mass
	}

	fn pos(&self) -> VecN<T, N> {
		self.pos
	}

	fn add_mass(&mut self, mass: T, pos: VecN<T, N>) {
		let new_mass = self.mass + mass;
		self.pos = (self.pos * self.mass + pos * mass) / new_mass;
		self.mass = new_mass;
	}

	fn custom(&self) -> C {
		self.custom
	}
}

pub(crate) fn apply_repulsion_2d<T: Coord + Send + Sync, E, O: Nodes<T, 2, Id>, Id: Sync>(
	layout: &mut Layout<T, 2, E, O, Id>,
) {
	let mut nodes_iter = layout.nodes.iter_nodes();
	let Some(first_node) = nodes_iter.next() else {
		return;
	};
	let (mut min_pos, mut max_pos) = (first_node.pos, first_node.pos);
	for Node { pos, .. } in nodes_iter {
		for ((min, max), val) in min_pos.iter_mut().zip(max_pos.iter_mut()).zip(pos.iter()) {
			if val < min {
				*min = *val;
			} else if val > max {
				*max = *val;
			}
		}
	}

	let mut bump = layout.bump.lock();
	let mut tree = crate::trees::Tree::<
		crate::trees::Node2<T, NodeBody2<T, ()>>,
		T,
		(),
		NodeBody2<T, ()>,
		2,
	>::from_bump(&mut bump);
	let mut root = tree.new_root((min_pos, max_pos));

	for node in layout.nodes.iter_nodes() {
		root.add_body(NodeBody2 {
			pos: node.pos,
			mass: node.mass + T::one(),
			custom: (),
		});
	}

	let kr = layout.settings.kr;
	let theta = layout.settings.theta;

	let f1 = |bb, b, mm, d, _| (bb - b) * mm / d;
	let f2 = |bb, b, mm, d, _, _| (bb - b) * mm / d;
	layout.nodes.par_iter_nodes_mut().for_each(|node| {
		node.speed -= root.apply(node.pos, theta, (), &f1, &f2) * kr * (node.mass + T::one());
	});
	std::mem::drop(root);
	tree.clear();
}

pub(crate) fn apply_repulsion_2d_po<T: Coord + Send + Sync, E, O: Nodes<T, 2, Id>, Id: Sync>(
	layout: &mut Layout<T, 2, E, O, Id>,
) {
	let mut kr_prime = layout.settings.prevent_overlapping.unwrap();

	let mut nodes_iter = layout.nodes.iter_nodes();
	let Some(first_node) = nodes_iter.next() else {
		return;
	};
	let (mut min_pos, mut max_pos) = (first_node.pos, first_node.pos);
	for Node { pos, .. } in nodes_iter {
		for ((min, max), val) in min_pos.iter_mut().zip(max_pos.iter_mut()).zip(pos.iter()) {
			if val < min {
				*min = *val;
			} else if val > max {
				*max = *val;
			}
		}
	}

	let mut bump = layout.bump.lock();
	let mut tree = crate::trees::Tree::<
		crate::trees::Node2<T, NodeBody2<T, T>>,
		T,
		T,
		NodeBody2<T, T>,
		2,
	>::from_bump(&mut bump);
	let mut root = tree.new_root((min_pos, max_pos));

	for node in layout.nodes.iter_nodes() {
		root.add_body(NodeBody2 {
			pos: node.pos,
			mass: node.mass + T::one(),
			custom: node.size,
		});
	}

	let kr = layout.settings.kr;
	kr_prime /= kr;
	let theta = layout.settings.theta;

	let f1 = |bb, b, mm, d, _| (bb - b) * mm / d;
	let f2 = |bb, b, mm, d, s1, s2| {
		let d_prime: T = d - s1 - s2;
		if d_prime.is_positive() {
			(bb - b) * mm / d_prime
		} else {
			(bb - b) * mm * kr_prime
		}
	};
	layout.nodes.par_iter_nodes_mut().for_each(|node| {
		node.speed -=
			root.apply(node.pos, theta, node.size, &f1, &f2) * kr * (node.mass + T::one());
	});
	std::mem::drop(root);
	tree.clear();
}

pub(crate) fn apply_repulsion_3d<T: Coord + Send + Sync, E, O: Nodes<T, 3, Id>, Id: Sync>(
	layout: &mut Layout<T, 3, E, O, Id>,
) {
	let mut nodes_iter = layout.nodes.iter_nodes();
	let Some(first_node) = nodes_iter.next() else {
		return;
	};
	let (mut min_pos, mut max_pos) = (first_node.pos, first_node.pos);
	for Node { pos, .. } in nodes_iter {
		for ((min, max), val) in min_pos.iter_mut().zip(max_pos.iter_mut()).zip(pos.iter()) {
			if val < min {
				*min = *val;
			} else if val > max {
				*max = *val;
			}
		}
	}

	let mut bump = layout.bump.lock();
	let mut tree = crate::trees::Tree::<
		crate::trees::Node3<T, NodeBody3<T, ()>>,
		T,
		(),
		NodeBody3<T, ()>,
		3,
	>::from_bump(&mut bump);
	let mut root = tree.new_root((min_pos, max_pos));

	for node in layout.nodes.iter_nodes() {
		root.add_body(NodeBody3 {
			pos: node.pos,
			mass: node.mass + T::one(),
			custom: (),
		});
	}

	let kr = layout.settings.kr;
	let theta = layout.settings.theta;

	layout.nodes.par_iter_nodes_mut().for_each(|node| {
		node.speed -= root.apply(
			node.pos,
			theta,
			(),
			&|bb, b, mm, d, _| (bb - b) * mm / d,
			&|bb, b, mm, d, _, _| (bb - b) * mm / d,
		) * kr * (node.mass + T::one());
	});
	std::mem::drop(root);
	tree.clear();
}

pub(crate) fn apply_repulsion_3d_po<T: Coord + Send + Sync, E, O: Nodes<T, 3, Id>, Id: Sync>(
	layout: &mut Layout<T, 3, E, O, Id>,
) {
	let mut kr_prime = layout.settings.prevent_overlapping.unwrap();

	let mut nodes_iter = layout.nodes.iter_nodes();
	let Some(first_node) = nodes_iter.next() else {
		return;
	};
	let (mut min_pos, mut max_pos) = (first_node.pos, first_node.pos);
	for Node { pos, .. } in nodes_iter {
		for ((min, max), val) in min_pos.iter_mut().zip(max_pos.iter_mut()).zip(pos.iter()) {
			if val < min {
				*min = *val;
			} else if val > max {
				*max = *val;
			}
		}
	}

	let mut bump = layout.bump.lock();
	let mut tree = crate::trees::Tree::<
		crate::trees::Node3<T, NodeBody3<T, T>>,
		T,
		T,
		NodeBody3<T, T>,
		3,
	>::from_bump(&mut bump);
	let mut root = tree.new_root((min_pos, max_pos));

	for node in layout.nodes.iter_nodes() {
		root.add_body(NodeBody3 {
			pos: node.pos,
			mass: node.mass + T::one(),
			custom: node.size,
		});
	}

	let kr = layout.settings.kr;
	kr_prime /= kr;
	let theta = layout.settings.theta;

	let f1 = |bb, b, mm, d, _| (bb - b) * mm / d;
	let f2 = |bb, b, mm, d, s1, s2| {
		let d_prime: T = d - s1 - s2;
		if d_prime.is_positive() {
			(bb - b) * mm / d_prime
		} else {
			(bb - b) * mm * kr_prime
		}
	};
	layout.nodes.par_iter_nodes_mut().for_each(|node| {
		node.speed -=
			root.apply(node.pos, theta, node.size, &f1, &f2) * kr * (node.mass + T::one());
	});
	std::mem::drop(root);
	tree.clear();
}