forceatlas2 0.8.0

use crate::{
	edge::*,
	forces::{Attraction, Gravity, Repulsion},
	node::*,
	util::*,
};

use num_traits::Zero;
use rayon::iter::ParallelIterator;
use std::marker::PhantomData;

/// Settings for the graph layout
#[derive(Clone, Debug)]
pub struct Settings<T> {
	/// Precision setting for Barnes-Hut computation
	///
	/// Must be in `(0.0..1.0)`. `0.0` is accurate and slow, `1.0` is unaccurate and fast.
	/// Default is `0.5`.
	pub theta: T,
	/// Attraction coefficient
	pub ka: T,
	/// Gravity coefficient
	pub kg: T,
	/// Repulsion coefficient
	pub kr: T,
	/// Logarithmic attraction
	pub lin_log: bool,
	/// Prevent node overlapping for a prettier graph.
	///
	/// Value is `kr_prime`.
	/// Requires `layout.sizes` to be `Some`.
	/// `kr_prime` is arbitrarily set to `100.0` in Gephi implementation.
	pub prevent_overlapping: Option<T>,
	/// Speed factor
	pub speed: T,
	/// Gravity does not decrease with distance, resulting in a more compact graph.
	pub strong_gravity: bool,
}

impl<T: Coord> Default for Settings<T> {
	fn default() -> Self {
		Self {
			theta: T::one() / (T::one() + T::one()),
			ka: T::one(),
			kg: T::one(),
			kr: T::one(),
			lin_log: false,
			prevent_overlapping: None,
			speed: T::from(0.01).unwrap_or_else(T::one),
			strong_gravity: false,
		}
	}
}

impl<T: Coord> Settings<T> {
	/// Check whether the settings are valid
	pub fn check(&self) -> bool {
		self.theta >= T::zero() && self.theta <= T::one()
	}
}

/// Graph spatialization layout
pub struct Layout<T, const N: usize, E = EdgeVec<T, usize>, O = NodeVec<T, N>, Id = usize> {
	/// Graph vertices positioned in space
	pub nodes: O,
	/// Graph edges (undirected)
	pub edges: E,
	// Mutex needed here for Layout to be Sync
	pub(crate) bump: parking_lot::Mutex<bumpalo::Bump>,
	pub(crate) fn_attraction: fn(&mut Self),
	pub(crate) fn_gravity: fn(&mut Self),
	pub(crate) fn_repulsion: fn(&mut Self),
	pub(crate) settings: Settings<T>,
	pub(crate) global_speed: T,
	pub(crate) _p: PhantomData<Id>,
}

impl<T: Coord + Send + Sync, const N: usize, E, O, Id> Default for Layout<T, N, E, O, Id>
where
	Layout<T, N, E, O, Id>:
		Attraction<T, N, E, O, Id> + Gravity<T, N, E, O, Id> + Repulsion<T, N, E, O, Id>,
	E: Default + Edges<T, Id>,
	O: Default + Nodes<T, N, Id>,
{
	fn default() -> Self {
		let settings = Default::default();
		Self {
			edges: E::default(),
			nodes: O::default(),
			bump: parking_lot::Mutex::new(bumpalo::Bump::new()),
			fn_attraction: Self::choose_attraction(&settings),
			fn_gravity: Self::choose_gravity(&settings),
			fn_repulsion: Self::choose_repulsion(&settings),
			settings,
			global_speed: T::one(),
			_p: Default::default(),
		}
	}
}

impl<T: Coord + Send + Sync, const N: usize, E, O, Id: Sync> Layout<T, N, E, O, Id>
where
	Layout<T, N, E, O, Id>:
		Attraction<T, N, E, O, Id> + Gravity<T, N, E, O, Id> + Repulsion<T, N, E, O, Id>,
	E: Edges<T, Id>,
	O: Nodes<T, N, Id>,
{
	/// Create an empty layout from settings
	pub fn empty(settings: Settings<T>) -> Self
	where
		E: Default,
		O: Default,
	{
		assert!(settings.check());
		Self {
			edges: E::default(),
			nodes: O::default(),
			bump: parking_lot::Mutex::new(bumpalo::Bump::new()),
			fn_attraction: Self::choose_attraction(&settings),
			fn_gravity: Self::choose_gravity(&settings),
			fn_repulsion: Self::choose_repulsion(&settings),
			settings,
			global_speed: T::one(),
			_p: Default::default(),
		}
	}

	/// Create a layout from an existing spatial graph
	#[cfg(feature = "rand")]
	pub fn from_abstract<R: rand::Rng>(
		settings: Settings<T>,
		nodes: impl Iterator<Item = (Id, AbstractNode<T>)>,
		edges: E,
		rng: &mut R,
	) -> Self
	where
		O: Default + BuildableNodes<T, N, Id>,
		rand::distributions::Standard: rand::distributions::Distribution<T>,
		T: rand::distributions::uniform::SampleUniform,
	{
		assert!(settings.check());
		let mut concrete_nodes = O::default();
		nodes.for_each(|(id, node)| {
			concrete_nodes.add_node(
				id,
				Node {
					pos: sample_unit_cube(rng),
					speed: Zero::zero(),
					old_speed: Zero::zero(),
					mass: node.mass,
					size: node.size,
				},
			)
		});
		Self {
			bump: parking_lot::Mutex::new(bumpalo::Bump::with_capacity(
				(concrete_nodes.len()
					+ 4 * (concrete_nodes.len().checked_ilog2().unwrap_or(0) as usize + 1))
					* std::mem::size_of::<
						crate::trees::NodeN<T, crate::forces::repulsion::NodeBodyN<T, T, N>, N, 1>,
					>(),
			)),
			edges,
			nodes: concrete_nodes,
			fn_attraction: Self::choose_attraction(&settings),
			fn_gravity: Self::choose_gravity(&settings),
			fn_repulsion: Self::choose_repulsion(&settings),
			settings,
			global_speed: T::one(),
			_p: Default::default(),
		}
	}

	/// Create a layout from an existing spatial graph
	pub fn from_positioned(settings: Settings<T>, nodes: O, edges: E) -> Self {
		assert!(settings.check());
		Self {
			bump: parking_lot::Mutex::new(bumpalo::Bump::with_capacity(
				(nodes.len() + 4 * (nodes.len().checked_ilog2().unwrap_or(0) as usize + 1))
					* std::mem::size_of::<
						crate::trees::NodeN<T, crate::forces::repulsion::NodeBodyN<T, T, N>, N, 1>,
					>(),
			)),
			edges,
			nodes,
			fn_attraction: Self::choose_attraction(&settings),
			fn_gravity: Self::choose_gravity(&settings),
			fn_repulsion: Self::choose_repulsion(&settings),
			settings,
			global_speed: T::one(),
			_p: Default::default(),
		}
	}

	/// Get layout's settings
	pub fn get_settings(&self) -> &Settings<T> {
		&self.settings
	}

	/// Change layout settings
	pub fn set_settings(&mut self, settings: Settings<T>) {
		assert!(settings.check());
		self.fn_attraction = Self::choose_attraction(&settings);
		self.fn_gravity = Self::choose_gravity(&settings);
		self.fn_repulsion = Self::choose_repulsion(&settings);
		self.settings = settings;
	}

	/// Compute an iteration of ForceAtlas2
	///
	/// Returns `(global_swinging, global_traction, max_traction)`.
	///
	/// * `global_swinging` measures how much the nodes change direction between two iterations.
	///   When the graph becomes stable, some nodes may swing or jiggle around a fixed position.
	///   This value helps detecting when this happens.
	/// * `global_traction` measures how much the nodes have a steady movement.
	///   When the graph becomes stable, this value should be low.
	///
	/// Depending on your application, parameters and graphs, you can empirically find thresholds for these values
	/// to detect when the simulation must be stopped.
	///
	/// Note that the graph may have angular momentum.
	pub fn iteration(&mut self) -> (T, T, T) {
		self.apply_attraction();
		self.apply_repulsion();
		self.apply_gravity();
		self.apply_forces()
	}

	/// Get the layout's global speed
	pub fn global_speed(&self) -> T {
		self.global_speed
	}

	fn apply_attraction(&mut self) {
		(self.fn_attraction)(self)
	}

	fn apply_gravity(&mut self) {
		(self.fn_gravity)(self)
	}

	fn apply_repulsion(&mut self) {
		(self.fn_repulsion)(self)
	}

	fn apply_forces(&mut self) -> (T, T, T) {
		let two = T::one() + T::one();
		let (global_swinging, global_traction, max_traction) = self
			.nodes
			.par_iter_nodes_mut()
			.map(|node| {
				// Swinging measures the node's instability
				let swinging: T = node
					.speed
					.iter()
					.zip(node.old_speed.iter())
					.map(|(s, old_s)| (*s - *old_s).powi(2))
					.sum::<T>()
					.sqrt();
				// Traction measures the node's movement stability
				// i.e. it will be higher when the node's direction is steady
				let traction: T = node
					.speed
					.iter()
					.zip(node.old_speed.iter())
					.map(|(s, old_s)| (*s + *old_s).powi(2))
					.sum::<T>()
					.sqrt() / two;

				// Reduce local speed when there is too much swinging
				//let f = self.global_speed / (self.global_speed * swinging.sqrt() + T::one()) * self.settings.speed;
				let f = traction.ln_1p() / (swinging.sqrt() + T::one()) * self.settings.speed;

				node.pos
					.iter_mut()
					.zip(node.speed.iter_mut())
					.for_each(|(pos, speed)| {
						*pos += *speed * f;
					});

				node.old_speed = node.speed;
				node.speed = Zero::zero();

				(swinging, traction, traction)
			})
			.reduce(
				|| (Zero::zero(), Zero::zero(), Zero::zero()),
				|a, b| (a.0 + b.0, a.1 + b.1, a.2.max(b.2)),
			);
		self.global_speed = (global_traction / global_swinging).min(two);
		(global_swinging, global_traction, max_traction)
	}
}

#[cfg(test)]
mod test {
	use super::*;

	fn is_send_sync<T: Send + Sync>() {}

	#[test]
	fn layout_send_sync() {
		is_send_sync::<Layout<f32, 2>>()
	}

	#[test]
	fn test_forces() {
		let mut layout = Layout::<f64, 2>::from_positioned(
			Settings::default(),
			vec![
				Node {
					pos: VecN([-2.0, -2.0]),
					..Default::default()
				},
				Node {
					pos: VecN([1.0, 2.0]),
					..Default::default()
				},
			],
			vec![((0, 1), 1.0)],
		);

		layout.apply_attraction();

		let speed_1 = dbg!(layout.nodes[0].speed);
		let speed_2 = dbg!(layout.nodes[1].speed);

		assert!(speed_1[0] > 0.0);
		assert!(speed_1[1] > 0.0);
		assert!(speed_2[0] < 0.0);
		assert!(speed_2[1] < 0.0);
		assert_eq!(speed_1[0], 3.0);
		assert_eq!(speed_1[1], 4.0);
		assert_eq!(speed_2[0], -3.0);
		assert_eq!(speed_2[1], -4.0);

		for node in layout.nodes.iter_nodes_mut() {
			node.speed = Zero::zero();
		}
		layout.apply_repulsion();

		let speed_1 = dbg!(layout.nodes[0].speed);
		let speed_2 = dbg!(layout.nodes[1].speed);

		assert!(speed_1[0] < 0.0);
		assert!(speed_1[1] < 0.0);
		assert!(speed_2[0] > 0.0);
		assert!(speed_2[1] > 0.0);
		assert!(speed_1[0] > -10.0);
		assert!(speed_1[1] > -10.0);
		assert!(speed_2[0] < 10.0);
		assert!(speed_2[1] < 10.0);

		for node in layout.nodes.iter_nodes_mut() {
			node.speed = Zero::zero();
		}
		layout.apply_gravity();

		let speed_1 = dbg!(layout.nodes[0].speed);
		let speed_2 = dbg!(layout.nodes[1].speed);

		assert!(speed_1[0] > 0.0);
		assert!(speed_1[1] > 0.0);
		assert!(speed_2[0] < 0.0);
		assert!(speed_2[1] < 0.0);
	}

	#[test]
	fn test_convergence() {
		let mut layout = Layout::<f64, 2>::from_positioned(
			Settings {
				ka: 0.5,
				kg: 0.01,
				kr: 0.01,
				lin_log: false,
				prevent_overlapping: None,
				speed: 1.0,
				strong_gravity: false,
				theta: 0.5,
			},
			vec![
				Node {
					pos: VecN([-1.1, -1.0]),
					..Default::default()
				},
				Node {
					pos: VecN([0.0, 0.0]),
					..Default::default()
				},
				Node {
					pos: VecN([1.0, 1.0]),
					..Default::default()
				},
			],
			vec![((0, 1), 1.0), ((1, 2), 1.0)],
		);

		for _ in 0..10 {
			println!("new iteration");
			layout.apply_attraction();
			layout.apply_repulsion();
			layout.apply_gravity();
			layout.apply_forces();

			let point_1 = layout.nodes[0].pos;
			let point_2 = layout.nodes[1].pos;
			dbg!(((point_2[0] - point_1[0]).powi(2) + (point_2[1] - point_1[1]).powi(2)).sqrt());
		}
	}

	#[test]
	fn test_convergence_po() {
		let mut layout = Layout::<f64, 2>::from_positioned(
			Settings {
				ka: 0.5,
				kg: 0.01,
				kr: 0.01,
				lin_log: false,
				prevent_overlapping: Some(100.),
				speed: 1.0,
				strong_gravity: false,
				theta: 0.5,
			},
			vec![
				Node {
					pos: VecN([-1.1, -1.0]),
					size: 1.0,
					..Default::default()
				},
				Node {
					pos: VecN([0.0, 0.0]),
					size: 5.0,
					..Default::default()
				},
				Node {
					pos: VecN([1.0, 1.0]),
					size: 1.0,
					..Default::default()
				},
			],
			vec![((0, 1), 1.0), ((1, 2), 1.0)],
		);

		for _ in 0..10 {
			println!("new iteration");
			layout.apply_attraction();
			layout.apply_repulsion();
			layout.apply_gravity();
			layout.apply_forces();

			let point_1 = layout.nodes[0].pos;
			let point_2 = layout.nodes[1].pos;
			dbg!(((point_2[0] - point_1[0]).powi(2) + (point_2[1] - point_1[1]).powi(2)).sqrt());
		}
	}
}