Parallel tree

Cargo.toml

@@ -7,7 +7,9 @@ battleship = []
 name = "bempp-octree"
 version = "0.0.1-dev"
 edition = "2021"
-authors = ["Srinath Kailasa <[email protected]>, Timo Betcke <[email protected]>"]
+authors = [
+    "Srinath Kailasa <[email protected]>, Timo Betcke <[email protected]>",
+]
 description = "A library to create Octrees"
 license = "BSD-3-Clause"
 homepage = "https://github.com/bempp/octree"
@@ -22,14 +24,15 @@ crate-type = ["cdylib", "lib"]
 
 [dependencies]
 itertools = "0.13.*"
-rand = "0.8.5"
-bytemuck = "1.*"
+rand = { version = "0.8.5", features = ["alloc"] }
+rand_chacha = "0.3.*"
+num = "0.4.*"
 vtkio = "0.6.*"
-mpi = {version = "0.8.*", features = ["derive", "user-operations"] }
+mpi = { version = "0.8.*", features = ["derive", "user-operations"] }
 
 [profile.release]
 debug = 1
- 
+
 [dev-dependencies]
 rand_distr = "0.4.3"
 #criterion = { version = "0.5.*", features = ["html_reports"]}

examples/battleship.rs

@@ -4,7 +4,7 @@
 use std::time::Instant;
 
 #[cfg(feature = "battleship")]
-use bempp_octree::octree::Octree;
+use bempp_octree::serial::Octree;
 #[cfg(feature = "battleship")]
 use vtkio::model::*;
 

examples/mpi_complete_tree.rs

@@ -0,0 +1,95 @@
+//! Test the computation of a complete octree.
+
+use bempp_octree::{
+    morton::MortonKey,
+    octree::{is_complete_linear_and_balanced, KeyType, Octree},
+    tools::{gather_to_all, generate_random_points},
+};
+use itertools::Itertools;
+use mpi::traits::Communicator;
+use rand::prelude::*;
+use rand_chacha::ChaCha8Rng;
+
+pub fn main() {
+    // Initialise MPI
+    let universe = mpi::initialize().unwrap();
+
+    // Get the world communicator
+    let comm = universe.world();
+
+    // Initialise a seeded Rng.
+    let mut rng = ChaCha8Rng::seed_from_u64(comm.rank() as u64);
+
+    // Create `npoints` per rank.
+    let npoints = 10000;
+
+    // Generate random points.
+
+    let points = generate_random_points(npoints, &mut rng, &comm);
+
+    let tree = Octree::new(&points, 15, 50, &comm);
+
+    // We now check that each node of the tree has all its neighbors available.
+
+    let leaf_tree = tree.leaf_tree();
+    let all_keys = tree.all_keys();
+
+    assert!(is_complete_linear_and_balanced(leaf_tree, &comm));
+    for &key in leaf_tree {
+        // We only check interior keys. Leaf keys may not have a neighbor
+        // on the same level.
+        let mut parent = key.parent();
+        while parent.level() > 0 {
+            // Check that the key itself is there.
+            assert!(all_keys.contains_key(&key));
+            // Check that all its neighbours are there.
+            for neighbor in parent.neighbours().iter().filter(|&key| key.is_valid()) {
+                assert!(all_keys.contains_key(neighbor));
+            }
+            parent = parent.parent();
+            // Check that the parent is there.
+            assert!(all_keys.contains_key(&parent));
+        }
+    }
+
+    // At the end check that the root of the tree is also contained.
+    assert!(all_keys.contains_key(&MortonKey::root()));
+
+    // Count the number of ghosts on each rank
+    // Count the number of global keys on each rank.
+
+    // Assert that all ghosts are from a different rank and count them.
+
+    let nghosts = all_keys
+        .iter()
+        .filter_map(|(_, &value)| {
+            if let KeyType::Ghost(rank) = value {
+                assert!(rank != comm.size() as usize);
+                Some(rank)
+            } else {
+                None
+            }
+        })
+        .count();
+
+    if comm.size() == 1 {
+        assert_eq!(nghosts, 0);
+    } else {
+        assert!(nghosts > 0);
+    }
+
+    let nglobal = all_keys
+        .iter()
+        .filter(|(_, &value)| matches!(value, KeyType::Global))
+        .count();
+
+    // Assert that all globals across all ranks have the same count.
+
+    let nglobals = gather_to_all(std::slice::from_ref(&nglobal), &comm);
+
+    assert_eq!(nglobals.iter().unique().count(), 1);
+
+    if comm.rank() == 0 {
+        println!("Distributed tree is complete and linear.");
+    }
+}

examples/mpi_cumsum.rs

@@ -0,0 +1,62 @@
+//! Test the computation of a global bounding box across MPI ranks.
+
+use bempp_octree::tools::{gather_to_root, global_inclusive_cumsum};
+use itertools::{izip, Itertools};
+use mpi::traits::Communicator;
+use rand::prelude::*;
+use rand_chacha::ChaCha8Rng;
+
+pub fn main() {
+    // Initialise MPI
+    let universe = mpi::initialize().unwrap();
+
+    // Get the world communicator
+    let comm = universe.world();
+
+    // Initialise a seeded Rng.
+    let mut rng = ChaCha8Rng::seed_from_u64(comm.rank() as u64);
+
+    // Create `npoints` per rank.
+    let nelems = 10;
+
+    // Generate random numbers
+
+    let mut elems = Vec::<usize>::with_capacity(nelems);
+
+    for _ in 0..nelems {
+        elems.push(rng.gen_range(0..100));
+    }
+
+    // Compute the cumulative sum.
+
+    let global_cum_sum = global_inclusive_cumsum(&elems, &comm);
+
+    // Copy array to root and compare with inclusive scan there.
+
+    if let (Some(cum_sum_root), Some(original_array)) = (
+        gather_to_root(&global_cum_sum, &comm),
+        gather_to_root(&elems, &comm),
+    ) {
+        // Scan on root
+
+        let expected_cum_sum = original_array
+            .iter()
+            .scan(0, |state, x| {
+                *state += *x;
+                Some(*state)
+            })
+            .collect_vec();
+
+        // Check that the first element is not modified (inclusive cumsum)
+        assert_eq!(
+            original_array.first().unwrap(),
+            cum_sum_root.first().unwrap()
+        );
+
+        for (actual, expected) in izip!(cum_sum_root.iter(), expected_cum_sum.iter()) {
+            assert_eq!(*actual, *expected);
+        }
+
+        println!("Cumulative sum computed.");
+    }
+}

examples/mpi_global_bounding_box.rs

@@ -0,0 +1,40 @@
+//! Test the computation of a global bounding box across MPI ranks.
+
+use bempp_octree::{
+    geometry::PhysicalBox,
+    octree::compute_global_bounding_box,
+    tools::{gather_to_root, generate_random_points},
+};
+use rand::prelude::*;
+use rand_chacha::ChaCha8Rng;
+
+pub fn main() {
+    // Initialise MPI
+    let universe = mpi::initialize().unwrap();
+
+    // Get the world communicator
+    let comm = universe.world();
+
+    // Initialise a seeded Rng.
+    let mut rng = ChaCha8Rng::seed_from_u64(2);
+
+    // Create `npoints` per rank.
+    let npoints = 10;
+
+    // Generate random points.
+
+    let points = generate_random_points(npoints, &mut rng, &comm);
+
+    // Compute the distributed bounding box.
+
+    let bounding_box = compute_global_bounding_box(&points, &comm);
+
+    // Copy all points to root and compare local bounding box there.
+
+    if let Some(points_root) = gather_to_root(&points, &comm) {
+        // Compute the bounding box on root.
+
+        let expected = PhysicalBox::from_points(&points_root);
+        assert_eq!(expected.coordinates(), bounding_box.coordinates());
+    }
+}

examples/parsort.rs

@@ -1,37 +1,25 @@
 //! Testing the hyksort component.
-use bempp_octree::parsort::{array_to_root, parsort};
-use itertools::Itertools;
+use bempp_octree::{
+    parsort::parsort,
+    tools::{generate_random_keys, is_sorted_array},
+};
 use mpi::traits::Communicator;
 use rand::prelude::*;
 
 pub fn main() {
     let universe = mpi::initialize().unwrap();
     let world = universe.world();
-    let rank = world.rank() as u64;
     let n_per_rank = 1000;
 
     let mut rng = rand::rngs::StdRng::seed_from_u64(0);
 
-    let mut arr = Vec::<u64>::new();
+    let keys = generate_random_keys(n_per_rank, &mut rng);
 
-    for _ in 0..n_per_rank {
-        arr.push(rng.gen());
-    }
-
-    // let splitters = get_splitters(&arr, &world, &mut rng);
-
-    // let bin_displs = get_bin_displacements(&arr, &splitters);
-
-    let arr = parsort(&arr, &world, &mut rng);
-    let arr = array_to_root(&arr, &world);
+    let arr = parsort(&keys, &world, &mut rng);
 
-    if rank == 0 {
-        let arr = arr.unwrap();
+    assert!(is_sorted_array(&arr, &world));
 
-        for (elem1, elem2) in arr.iter().tuple_windows() {
-            assert!(elem1 <= elem2);
-        }
-        println!("Sorted {} elements.", arr.len());
-        println!("Finished.");
+    if world.rank() == 0 {
+        println!("Array is sorted.");
     }
 }

src/geometry.rs

@@ -1,9 +1,33 @@
 //! Geometry information
 
-use bytemuck;
+use mpi::traits::Equivalence;
 
 use crate::constants::DEEPEST_LEVEL;
 
+/// Definition of a point.
+#[derive(Clone, Copy, Equivalence)]
+pub struct Point {
+    coords: [f64; 3],
+    global_id: usize,
+}
+
+impl Point {
+    /// Create a new point from coordinates and global id.
+    pub fn new(coords: [f64; 3], global_id: usize) -> Self {
+        Self { coords, global_id }
+    }
+
+    /// Return the coordintes of a point.
+    pub fn coords(&self) -> [f64; 3] {
+        self.coords
+    }
+
+    /// Return the global id of the point.
+    pub fn global_id(&self) -> usize {
+        self.global_id
+    }
+}
+
 /// A bounding box describes geometry in which an Octree lives.
 pub struct PhysicalBox {
     coords: [f64; 6],
@@ -18,11 +42,7 @@ impl PhysicalBox {
     }
 
     /// Give a slice of points. Compute an associated bounding box.
-    pub fn from_points(points: &[f64]) -> PhysicalBox {
-        assert_eq!(points.len() % 3, 0);
-
-        let points: &[[f64; 3]] = bytemuck::cast_slice(points);
-
+    pub fn from_points(points: &[Point]) -> PhysicalBox {
         let mut xmin = f64::MAX;
         let mut xmax = f64::MIN;
 
@@ -33,9 +53,9 @@ impl PhysicalBox {
         let mut zmax = f64::MIN;
 
         for point in points {
-            let x = point[0];
-            let y = point[1];
-            let z = point[2];
+            let x = point.coords()[0];
+            let y = point.coords()[1];
+            let z = point.coords()[2];
 
             xmin = f64::min(xmin, x);
             xmax = f64::max(xmax, x);

src/lib.rs

@@ -7,4 +7,6 @@ pub mod geometry;
 pub mod morton;
 pub mod octree;
 pub mod parsort;
+pub mod serial;
+pub mod tools;
 pub mod types;