diff --git a/Cargo.toml b/Cargo.toml index 81519282..58fdd422 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,6 +36,9 @@ rayon = "1.9" rlst = { git = "https://github.com/linalg-rs/rlst.git", features = ["mpi"] } green-kernels = { git = "https://github.com/bempp/green-kernels.git", features = ["mpi"] } # c-api-tools = { version = "0.1.0" } +kifmm = { git = "https://github.com/bempp/kifmm.git", features = ["mpi"] } +bempp-distributed-tools = { git = "https://github.com/bempp/distributed_tools.git"} + [dev-dependencies] approx = "0.5" diff --git a/src/greens_function_evaluators.rs b/src/greens_function_evaluators.rs index 9b4eb1c1..183bee19 100644 --- a/src/greens_function_evaluators.rs +++ b/src/greens_function_evaluators.rs @@ -1,3 +1,4 @@ //! Interfaces to Green's function evaluators. pub mod dense_evaluator; +pub mod kifmm_evaluator; diff --git a/src/greens_function_evaluators/kifmm_evaluator.rs b/src/greens_function_evaluators/kifmm_evaluator.rs new file mode 100644 index 00000000..ecd25e03 --- /dev/null +++ b/src/greens_function_evaluators/kifmm_evaluator.rs @@ -0,0 +1,281 @@ +//! Interface to kifmm library + +use std::cell::RefCell; + +use bempp_distributed_tools::{self, permutation::DataPermutation}; + +use green_kernels::{laplace_3d::Laplace3dKernel, types::GreenKernelEvalType}; +use kifmm::{ + traits::{ + fftw::{Dft, DftType}, + field::{ + SourceToTargetTranslation, SourceToTargetTranslationMetadata, SourceTranslation, + TargetTranslation, + }, + fmm::{DataAccessMulti, EvaluateMulti}, + general::{ + multi_node::GhostExchange, + single_node::{AsComplex, Epsilon}, + }, + }, + tree::{Domain, SortKind}, + ChargeHandler, Evaluate, FftFieldTranslation, KiFmm, KiFmmMulti, MultiNodeBuilder, +}; +use mpi::{ + topology::SimpleCommunicator, + traits::{Communicator, Equivalence}, +}; +use num::Float; +use rayon::range; +use rlst::{ + operator::interface::DistributedArrayVectorSpace, rlst_dynamic_array1, AsApply, Element, + IndexLayout, MatrixSvd, OperatorBase, RawAccess, RawAccessMut, RlstScalar, +}; + +pub struct KiFmmEvaluator< + 'a, + C: Communicator, + T: RlstScalar + Equivalence, + SourceLayout: IndexLayout, + TargetLayout: IndexLayout, +> where + T::Real: Equivalence, + T: DftType::ComplexType>, + T: Dft + AsComplex + Epsilon + MatrixSvd + Float, + KiFmmMulti, FftFieldTranslation>: SourceToTargetTranslationMetadata, +{ + sources: Vec, + targets: Vec, + eval_mode: GreenKernelEvalType, + domain_space: &'a DistributedArrayVectorSpace<'a, SourceLayout, T>, + range_space: &'a DistributedArrayVectorSpace<'a, TargetLayout, T>, + source_permutation: DataPermutation<'a, SourceLayout>, + target_permutation: DataPermutation<'a, TargetLayout>, + n_permuted_sources: usize, + n_permuted_targets: usize, + simple_comm: SimpleCommunicator, + fmm: RefCell, FftFieldTranslation>>, +} + +impl< + 'a, + C: Communicator, + T: RlstScalar + Equivalence, + SourceLayout: IndexLayout, + TargetLayout: IndexLayout, + > std::fmt::Debug for KiFmmEvaluator<'a, C, T, SourceLayout, TargetLayout> +where + T::Real: Equivalence, + T: DftType::ComplexType>, + T: Dft + AsComplex + Epsilon + MatrixSvd + Float, + KiFmmMulti, FftFieldTranslation>: SourceToTargetTranslationMetadata, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "KiFmmEvaluator with {} sources and {} targets", + self.domain_space.index_layout().number_of_global_indices(), + self.range_space.index_layout().number_of_global_indices() + ) + } +} + +impl< + 'a, + C: Communicator, + T: RlstScalar + Equivalence, + SourceLayout: IndexLayout, + TargetLayout: IndexLayout, + > KiFmmEvaluator<'a, C, T, SourceLayout, TargetLayout> +where + T::Real: Equivalence, + T: DftType::ComplexType>, + T: Dft + AsComplex + Epsilon + MatrixSvd + Float, + KiFmmMulti, FftFieldTranslation>: SourceToTargetTranslationMetadata, +{ + /// Instantiate a new KiFmm Evaluator. + pub fn new( + sources: &[T::Real], + targets: &[T::Real], + eval_mode: GreenKernelEvalType, + local_tree_depth: usize, + global_tree_depth: usize, + expansion_order: usize, + domain_space: &'a DistributedArrayVectorSpace<'a, SourceLayout, T>, + range_space: &'a DistributedArrayVectorSpace<'a, TargetLayout, T>, + ) -> Self { + // We want that both layouts have the same communicator. + assert!(std::ptr::addr_eq(domain_space.comm(), range_space.comm())); + + assert_eq!( + sources.len() % 3, + 0, + "Source vector length must be a multiple of 3." + ); + assert_eq!( + targets.len() % 3, + 0, + "Target vector length must be a multiple of 3." + ); + + // The length of the source vector must be 3 times the length of the local source indices. + assert_eq!( + sources.len(), + 3 * domain_space.index_layout().number_of_local_indices(), + "Number of sources ({}) does not match number of local indices ({}).", + sources.len() / 3, + domain_space.index_layout().number_of_local_indices(), + ); + + // The length of the target vector must be 3 times the length of the local target indices. + assert_eq!( + targets.len(), + 3 * range_space.index_layout().number_of_local_indices(), + "Number of targets ({}) does not match number of local indices ({}).", + targets.len() / 3, + range_space.index_layout().number_of_local_indices(), + ); + + let simple_comm = domain_space.index_layout().comm().duplicate(); + let cell = RefCell::new( + MultiNodeBuilder::new(false) + .tree( + &simple_comm, + sources, + targets, + local_tree_depth as u64, + global_tree_depth as u64, + true, + SortKind::Samplesort { n_samples: 100 }, + ) + .unwrap() + .parameters( + expansion_order, + Laplace3dKernel::::default(), + FftFieldTranslation::::new(None), + ) + .unwrap() + .build() + .unwrap(), + ); + + let source_indices = cell.borrow().tree.source_tree.global_indices.clone(); + let target_indices = cell.borrow().tree.target_tree.global_indices.clone(); + + let source_permutation = DataPermutation::new(domain_space.index_layout(), &source_indices); + let target_permutation = DataPermutation::new(range_space.index_layout(), &target_indices); + + KiFmmEvaluator { + sources: sources.to_vec(), + targets: targets.to_vec(), + eval_mode, + domain_space, + range_space, + source_permutation, + target_permutation, + n_permuted_sources: source_indices.len(), + n_permuted_targets: target_indices.len(), + simple_comm, + fmm: cell, + } + } +} + +impl< + 'a, + C: Communicator, + T: RlstScalar + Equivalence, + SourceLayout: IndexLayout, + TargetLayout: IndexLayout, + > OperatorBase for KiFmmEvaluator<'a, C, T, SourceLayout, TargetLayout> +where + T::Real: Equivalence, + T: DftType::ComplexType>, + T: Dft + AsComplex + Epsilon + MatrixSvd + Float, + KiFmmMulti, FftFieldTranslation>: SourceToTargetTranslationMetadata, +{ + type Domain = DistributedArrayVectorSpace<'a, SourceLayout, T>; + + type Range = DistributedArrayVectorSpace<'a, TargetLayout, T>; + + fn domain(&self) -> &Self::Domain { + self.domain_space + } + + fn range(&self) -> &Self::Range { + self.range_space + } +} + +impl< + 'a, + C: Communicator, + T: RlstScalar + Equivalence, + SourceLayout: IndexLayout, + TargetLayout: IndexLayout, + > AsApply for KiFmmEvaluator<'a, C, T, SourceLayout, TargetLayout> +where + T::Real: Equivalence, + T: DftType::ComplexType>, + T: Dft + AsComplex + Epsilon + MatrixSvd + Float, + KiFmmMulti, FftFieldTranslation>: SourceToTargetTranslationMetadata, + KiFmmMulti, FftFieldTranslation>: SourceToTargetTranslation, + KiFmm, FftFieldTranslation>: Evaluate, + KiFmmMulti, FftFieldTranslation>: SourceTranslation, + KiFmmMulti, FftFieldTranslation>: TargetTranslation, + KiFmmMulti, FftFieldTranslation>: DataAccessMulti, + KiFmmMulti, FftFieldTranslation>: GhostExchange, +{ + fn apply_extended( + &self, + alpha: ::F, + x: &::E, + beta: ::F, + y: &mut ::E, + ) -> rlst::RlstResult<()> { + let mut x_permuted = rlst_dynamic_array1![T, [self.n_permuted_sources]]; + let mut y_permuted = rlst_dynamic_array1![T, [self.n_permuted_targets]]; + + // This is the result vector of the FMM when permuted back into proper ordering. + let mut y_result = rlst_dynamic_array1![ + T, + [self.range_space.index_layout().number_of_local_indices()] + ]; + + // Now forward permute the source vector. + + self.source_permutation + .forward_permute(x.view().local().data(), x_permuted.data_mut()); + + // Multiply with the vector alpha + + x_permuted.scale_inplace(alpha); + + // We can now attach the charges to the kifmm. + + self.fmm + .borrow_mut() + .attach_charges_ordered(x_permuted.data()) + .unwrap(); + + // Now we can evaluate the FMM. + + self.fmm.borrow_mut().evaluate().unwrap(); + + // Save it into y_permuted. + y_permuted + .data_mut() + .copy_from_slice(self.fmm.borrow().potentials().unwrap().as_slice()); + + // Now permute back the result. + self.target_permutation + .backward_permute(y_permuted.data(), y_result.data_mut()); + + // Scale the vector y with beta before we add the result. + y.scale_inplace(beta); + // Now add the result. + y.view_mut().local_mut().sum_into(y_result.r()); + + Ok(()) + } +}