Add flop counter

kifmm/src/tree/helpers.rs

@@ -45,36 +45,76 @@ pub fn points_fixture<T: Float + RlstScalar + rand::distributions::uniform::Samp
     points
 }
 
-/// Points fixture for testing, uniformly samples on surface of a sphere of diameter 1.
-///
-/// # Arguments
-/// * `n_points` - The number of points to sample.
-/// * `min` - The minimum coordinate value along each axis.
-/// * `max` - The maximum coordinate value along each axis.
 pub fn points_fixture_sphere<T: RlstScalar + rand::distributions::uniform::SampleUniform>(
     n_points: usize,
 ) -> PointsMat<T> {
-    // Generate a set of randomly distributed points
-    let mut range = StdRng::seed_from_u64(0);
-    let pi = T::from(3.134159).unwrap();
+    use rand::distributions::Uniform;
+    use rand::rngs::StdRng;
+    use rand::SeedableRng;
+
+    // Seeded random number generator for reproducibility
+    let mut rng = StdRng::seed_from_u64(0);
+    let pi = T::from(std::f64::consts::PI).unwrap();
     let two = T::from(2.0).unwrap();
-    let half = T::from(0.5).unwrap();
+    let one = T::one();
 
-    let between = rand::distributions::Uniform::from(T::zero()..T::one());
+    // Uniform distributions for phi and cos(theta)
+    let phi_dist = Uniform::from(T::zero()..(two * pi));
+    let u_dist = Uniform::from(T::zero()..one);
 
+    // Initialize points array
     let mut points = rlst_dynamic_array2!(T, [3, n_points]);
-    let mut phi = rlst_dynamic_array2!(T, [n_points, 1]);
-    let mut theta = rlst_dynamic_array2!(T, [n_points, 1]);
 
     for i in 0..n_points {
-        phi[[i, 0]] = between.sample(&mut range) * two * pi;
-        theta[[i, 0]] = ((between.sample(&mut range) - half) * two).acos();
+        // Generate random phi and theta
+        let phi = phi_dist.sample(&mut rng);
+        let u = u_dist.sample(&mut rng);
+        let theta = (one - two * u).acos();
+
+        // Compute Cartesian coordinates
+        let x = theta.sin() * phi.cos();
+        let y = theta.sin() * phi.sin();
+        let z = theta.cos();
+
+        // Assign to points array
+        points[[0, i]] = x;
+        points[[1, i]] = y;
+        points[[2, i]] = z;
     }
 
+    points
+}
+
+pub fn points_fixture_oblate_spheroid<
+    T: RlstScalar + rand::distributions::uniform::SampleUniform,
+>(
+    n_points: usize,
+    a: T, // Semi-axis length along x- and y-axes
+    c: T, // Semi-axis length along the z-axis
+) -> PointsMat<T> {
+    use rand::distributions::Uniform;
+    use rand::rngs::StdRng;
+    use rand::SeedableRng;
+
+    let mut rng = StdRng::seed_from_u64(0); // Use a fixed seed for reproducibility
+    let pi = T::from(std::f64::consts::PI).unwrap();
+    let two = T::from(2.0).unwrap();
+    let one = T::one();
+
+    let phi_dist = Uniform::from(T::zero()..(two * pi)); // Azimuthal angle
+    let cos_theta_dist = Uniform::from(-one..one); // Cosine of polar angle
+
+    let mut points = rlst_dynamic_array2!(T, [3, n_points]);
+
     for i in 0..n_points {
-        points[[0, i]] = half * theta[[i, 0]].sin() * phi[[i, 0]].cos() + half;
-        points[[1, i]] = half * theta[[i, 0]].sin() * phi[[i, 0]].sin() + half;
-        points[[2, i]] = half * theta[[i, 0]].cos() + half;
+        let phi = phi_dist.sample(&mut rng); // Random azimuthal angle
+        let cos_theta = cos_theta_dist.sample(&mut rng); // Random cosine of polar angle
+        let theta = cos_theta.acos(); // Compute polar angle
+
+        // Parametric equations for the oblate spheroid
+        points[[0, i]] = a * theta.sin() * phi.cos(); // x-axis
+        points[[1, i]] = a * theta.sin() * phi.sin(); // y-axis
+        points[[2, i]] = c * theta.cos(); // z-axis
     }
 
     points

scripts/src/bin/count_flops_f32_amd.rs

@@ -0,0 +1,187 @@
+use core::panic;
+use std::time::Duration;
+
+use green_kernels::{laplace_3d::Laplace3dKernel, types::GreenKernelEvalType};
+use kifmm::fmm::types::FmmSvdMode;
+use kifmm::fmm::types::{BlasFieldTranslationSaRcmp, FftFieldTranslation, SingleNodeBuilder};
+use kifmm::traits::field::{SourceToTargetTranslation, TargetTranslation};
+use kifmm::traits::fmm::{DataAccess, Evaluate};
+use kifmm::traits::tree::{SingleFmmTree, SingleTree};
+use kifmm::tree::helpers::{points_fixture, points_fixture_oblate_spheroid, points_fixture_sphere};
+use rlst::{rlst_dynamic_array2, RawAccess, RawAccessMut};
+
+fn fft_f32() {
+    let n_points = 1000000;
+    let geometries = ["uniform", "sphere", "spheroid"];
+
+    // Tree depth
+    let depth_vec = vec![
+        [4, 5, 5], // 3 digits, for each geometry
+        [4, 5, 5], // 4 digits for each geometry
+    ];
+
+    // Expansion order
+    let e_vec = vec![
+        [3, 3, 3], // 3 digits, for each geometry
+        [4, 5, 5], // 4 digits for each geometry
+    ];
+
+    // Block size
+    let b_vec = vec![
+        [128, 64, 128], // 3 digits, for each geometry
+        [16, 64, 64], // 4 digits for each geometry
+    ];
+
+    let experiments = [3, 4]; // number of digits sought
+
+    let prune_empty = true;
+
+    for (i, digits) in experiments.iter().enumerate() {
+        for (j, &geometry) in geometries.iter().enumerate() {
+            let sources;
+            let targets;
+            if geometry == "uniform" {
+                sources = points_fixture::<f32>(n_points, Some(0.), Some(1.), None);
+                targets = points_fixture::<f32>(n_points, Some(0.), Some(1.), None);
+            } else if geometry == "sphere" {
+                sources = points_fixture_sphere(n_points);
+                targets = points_fixture_sphere(n_points);
+            } else if geometry == "spheroid" {
+                sources = points_fixture_oblate_spheroid(n_points, 1.0, 0.5);
+                targets = points_fixture_oblate_spheroid(n_points, 1.0, 0.5);
+            } else {
+                panic!("invalid geometry");
+            }
+
+            let charges = vec![1f32; n_points];
+
+            let depth = Some(depth_vec[i][j]);
+            let e = e_vec[i][j];
+            let b = b_vec[i][j];
+            let expansion_order = vec![e; depth.unwrap() as usize + 1];
+            let block_size = Some(b);
+
+            let mut fmm = SingleNodeBuilder::new(false)
+                .tree(sources.data(), targets.data(), None, depth, prune_empty)
+                .unwrap()
+                .parameters(
+                    &charges,
+                    &expansion_order,
+                    Laplace3dKernel::new(),
+                    GreenKernelEvalType::Value,
+                    FftFieldTranslation::new(block_size),
+                )
+                .unwrap()
+                .build()
+                .unwrap();
+
+            fmm.evaluate();
+
+            println!("FFT-M2L geometry {:?}, digits {:?}, flops {:?}", geometry, digits, fmm.nflops)
+
+        }
+    }
+}
+
+fn blas_f32() {
+
+    let n_points = 1000000;
+    let geometries = ["uniform", "sphere", "spheroid"];
+
+    // Tree depth
+    let depth_vec = vec![
+        [4, 4, 4], // 3 digits, for each geometry
+        [4, 5, 5], // 4 digits for each geometry
+    ];
+
+    // Expansion order
+    let e_vec = vec![
+        [3, 2, 3], // 3 digits, for each geometry
+        [3, 3, 4], // 4 digits for each geometry
+    ];
+
+    // SVD threshold
+    let svd_threshold_vec = vec![
+        [Some(0.00001), Some(0.00001), Some(0.001)],              // 3 digits
+        [None, None, None], // 4 digits
+    ];
+
+    let svd_mode_vec = vec![
+        [
+            FmmSvdMode::new(true, None, None, Some(10), None),
+            FmmSvdMode::new(true, None, None, Some(10), None),
+            FmmSvdMode::new(false, None, None, None, None),
+        ],
+        [
+            FmmSvdMode::new(false, None, None, None, None),
+            FmmSvdMode::new(true, None, None, Some(20), None),
+            FmmSvdMode::new(true, None, None, Some(5), None),
+        ],
+    ];
+
+    let surface_diff_vec = vec![
+        [None, Some(2), Some(2)], // 3 digits
+        [Some(1), Some(2), Some(1)],       // 4 digits
+    ];
+
+    let nvecs = vec![1];
+
+    let experiments = [3, 4]; // number of digits sought
+
+    let prune_empty = true;
+
+    for (i, digits) in experiments.iter().enumerate() {
+        for (j, &geometry) in geometries.iter().enumerate() {
+            for (k, nvec) in nvecs.iter().enumerate() {
+                let sources;
+                let targets;
+                if geometry == "uniform" {
+                    sources = points_fixture::<f32>(n_points, Some(0.), Some(1.), None);
+                    targets = points_fixture::<f32>(n_points, Some(0.), Some(1.), None);
+                } else if geometry == "sphere" {
+                    sources = points_fixture_sphere(n_points);
+                    targets = points_fixture_sphere(n_points);
+                } else if geometry == "spheroid" {
+                    sources = points_fixture_oblate_spheroid(n_points, 1.0, 0.5);
+                    targets = points_fixture_oblate_spheroid(n_points, 1.0, 0.5);
+                } else {
+                    panic!("invalid geometry");
+                }
+
+                let charges = vec![1f32; n_points * nvec];
+
+                let depth = Some(depth_vec[i][j]);
+                let e = e_vec[i][j];
+                let threshold = svd_threshold_vec[i][j];
+                let surface_diff = surface_diff_vec[i][j];
+                let svd_mode = svd_mode_vec[i][j];
+
+                let expansion_order = vec![e; depth.unwrap() as usize + 1];
+
+                let mut fmm = SingleNodeBuilder::new(false)
+                    .tree(sources.data(), targets.data(), None, depth, prune_empty)
+                    .unwrap()
+                    .parameters(
+                        &charges,
+                        &expansion_order,
+                        Laplace3dKernel::new(),
+                        GreenKernelEvalType::Value,
+                        BlasFieldTranslationSaRcmp::new(threshold, surface_diff, svd_mode),
+                    )
+                    .unwrap()
+                    .build()
+                    .unwrap();
+
+            fmm.evaluate();
+            println!("BLAS-M2L geometry {:?}, digits {:?}, flops {:?}", geometry, digits, fmm.nflops)
+            }
+        }
+    }
+}
+
+fn main() {
+
+    fft_f32();;
+
+    blas_f32();
+}