From 5f4f6c26c455c7d2ebd74db851e842dc8e94ec96 Mon Sep 17 00:00:00 2001 From: Timo Betcke Date: Thu, 2 Jan 2025 18:44:04 +0100 Subject: [PATCH] WIP: Evaluator tools --- benches/assembly_benchmark.rs | 3 +- examples/neighbour_evaluator.rs | 60 +++ src/boundary_assemblers.rs | 48 ++- src/evaluator_tools.rs | 248 ++++++++++-- src/function.rs | 378 ++++++++++++------ .../dense_evaluator.rs | 6 +- .../kifmm_evaluator.rs | 6 +- tests/test_space.rs | 31 +- 8 files changed, 592 insertions(+), 188 deletions(-) create mode 100644 examples/neighbour_evaluator.rs diff --git a/benches/assembly_benchmark.rs b/benches/assembly_benchmark.rs index 8a020f4f..847874cf 100644 --- a/benches/assembly_benchmark.rs +++ b/benches/assembly_benchmark.rs @@ -1,9 +1,8 @@ use bempp::boundary_assemblers::BoundaryAssemblerOptions; use bempp::function::FunctionSpace; -use bempp::function::FunctionSpaceTrait; +use bempp::function::LocalFunctionSpaceTrait; use bempp::laplace::assembler::single_layer; use criterion::{black_box, criterion_group, criterion_main, Criterion}; -use mpi; use ndelement::ciarlet::LagrangeElementFamily; use ndelement::types::{Continuity, ReferenceCellType}; diff --git a/examples/neighbour_evaluator.rs b/examples/neighbour_evaluator.rs new file mode 100644 index 00000000..9a1c071f --- /dev/null +++ b/examples/neighbour_evaluator.rs @@ -0,0 +1,60 @@ +//! Test the neighbour evaluator + +use bempp::evaluator_tools::NeighbourEvaluator; +use bempp_distributed_tools::IndexLayoutFromLocalCounts; +use green_kernels::laplace_3d::Laplace3dKernel; +use mpi::traits::Communicator; +use ndgrid::{ + traits::{Entity, Grid}, + types::Ownership, +}; +use rand::SeedableRng; +use rand_chacha::ChaCha8Rng; +use rlst::{ + operator::interface::DistributedArrayVectorSpace, rlst_dynamic_array2, AsApply, Element, + LinearSpace, OperatorBase, RawAccess, +}; + +fn main() { + let universe = mpi::initialize().unwrap(); + let world = universe.world(); + + let n_points = 5; + + let mut rng = ChaCha8Rng::seed_from_u64(world.rank() as u64); + + let mut points = rlst_dynamic_array2!(f64, [2, n_points]); + points.fill_from_equally_distributed(&mut rng); + + let grid = bempp::shapes::regular_sphere::(100, 1, &world); + + // Now get the active cells on the current process. + + let n_cells = grid + .entity_iter(2) + .filter(|e| matches!(e.ownership(), Ownership::Owned)) + .count(); + + let index_layout = IndexLayoutFromLocalCounts::new(n_cells * n_points, &world); + + let space = DistributedArrayVectorSpace::<_, f64>::new(&index_layout); + + let evaluator = NeighbourEvaluator::new( + points.data(), + Laplace3dKernel::default(), + &space, + &space, + &grid, + ); + + // Create an element in the space. + + let mut x = space.zero(); + let mut y = space.zero(); + + x.view_mut() + .local_mut() + .fill_from_equally_distributed(&mut rng); + + //let res = evaluator.apply(&x); +} diff --git a/src/boundary_assemblers.rs b/src/boundary_assemblers.rs index 854088ca..27d37a92 100644 --- a/src/boundary_assemblers.rs +++ b/src/boundary_assemblers.rs @@ -8,7 +8,7 @@ use crate::boundary_assemblers::cell_pair_assemblers::{ }; use crate::boundary_assemblers::helpers::KernelEvaluator; use crate::boundary_assemblers::helpers::{equal_grids, RawData2D, RlstArray, SparseMatrixData}; -use crate::function::FunctionSpaceTrait; +use crate::function::{FunctionSpaceTrait, LocalFunctionSpaceTrait}; use bempp_quadrature::duffy::{ quadrilateral_duffy, quadrilateral_triangle_duffy, triangle_duffy, triangle_quadrilateral_duffy, }; @@ -16,6 +16,7 @@ use bempp_quadrature::types::{CellToCellConnectivity, TestTrialNumericalQuadratu use green_kernels::traits::Kernel; use integrands::BoundaryIntegrand; use itertools::izip; +use mpi::traits::Communicator; use ndelement::quadrature::simplex_rule; use ndelement::reference_cell; use ndelement::traits::FiniteElement; @@ -119,13 +120,17 @@ impl<'o, T: RlstScalar + MatrixInverse, Integrand: BoundaryIntegrand, K: BoundaryAssembler<'o, T, Integrand, K> { /// Assemble the singular part into a CSR matrix. - pub fn assemble_singular + Sync>( + pub fn assemble_singular>( &self, trial_space: &Space, test_space: &Space, - ) -> CsrMatrix { + ) -> CsrMatrix + where + Space::LocalFunctionSpace: Sync, + { let shape = [test_space.global_size(), trial_space.global_size()]; - let sparse_matrix = self.assemble_singular_part(shape, trial_space, test_space); + let sparse_matrix = + self.assemble_singular_part(shape, trial_space.local_space(), test_space.local_space()); if sparse_matrix.data.is_empty() || sparse_matrix @@ -155,12 +160,15 @@ impl<'o, T: RlstScalar + MatrixInverse, Integrand: BoundaryIntegrand, K: } /// Assemble into a dense matrix. - pub fn assemble + Sync>( + pub fn assemble>( &self, trial_space: &Space, test_space: &Space, - ) -> DynamicArray { - if !trial_space.is_serial() || !test_space.is_serial() { + ) -> DynamicArray + where + Space::LocalFunctionSpace: Sync, + { + if trial_space.comm().size() > 1 || test_space.comm().size() > 1 { panic!("Dense assembly can only be used for function spaces stored in serial"); } @@ -173,17 +181,19 @@ impl<'o, T: RlstScalar + MatrixInverse, Integrand: BoundaryIntegrand, K: } /// Assemble into a dense matrix. - pub fn assemble_into_memory + Sync>( + pub fn assemble_into_memory>( &self, trial_space: &Space, test_space: &Space, output: &mut [T], - ) { + ) where + Space::LocalFunctionSpace: Sync, + { assert_eq!( output.len(), test_space.global_size() * trial_space.global_size() ); - if !trial_space.is_serial() || !test_space.is_serial() { + if trial_space.comm().size() > 1 || test_space.comm().size() > 1 { panic!("Dense assembly can only be used for function spaces stored in serial"); } @@ -197,13 +207,14 @@ impl<'o, T: RlstScalar + MatrixInverse, Integrand: BoundaryIntegrand, K: self.assemble_nonsingular_part( &output_raw, - trial_space, - test_space, + trial_space.local_space(), + test_space.local_space(), &trial_colouring, &test_colouring, ); - let sparse_matrix = self.assemble_singular_part(shape, trial_space, test_space); + let sparse_matrix = + self.assemble_singular_part(shape, trial_space.local_space(), test_space.local_space()); let data = sparse_matrix.data; let rows = sparse_matrix.rows; @@ -231,7 +242,7 @@ impl<'o, T: RlstScalar + MatrixInverse, Integrand: BoundaryIntegrand, K: } /// Assemble the singular contributions - fn assemble_singular_part + Sync>( + fn assemble_singular_part + Sync>( &self, shape: [usize; 2], trial_space: &Space, @@ -393,7 +404,7 @@ impl<'o, T: RlstScalar + MatrixInverse, Integrand: BoundaryIntegrand, K: } /// Assemble the non-singular contributions into a dense matrix - fn assemble_nonsingular_part + Sync>( + fn assemble_nonsingular_part + Sync>( &self, output: &RawData2D, trial_space: &Space, @@ -401,9 +412,6 @@ impl<'o, T: RlstScalar + MatrixInverse, Integrand: BoundaryIntegrand, K: trial_colouring: &HashMap>>, test_colouring: &HashMap>>, ) { - if !trial_space.is_serial() || !test_space.is_serial() { - panic!("Dense assembly can only be used for function spaces stored in serial"); - } if output.shape[0] != test_space.global_size() || output.shape[1] != trial_space.global_size() { @@ -609,7 +617,7 @@ where #[allow(clippy::too_many_arguments)] fn assemble_batch_singular< T: RlstScalar + MatrixInverse, - Space: FunctionSpaceTrait, + Space: LocalFunctionSpaceTrait + Sync, Integrand: BoundaryIntegrand, K: Kernel, >( @@ -688,7 +696,7 @@ fn assemble_batch_singular< #[allow(clippy::too_many_arguments)] fn assemble_batch_nonadjacent< T: RlstScalar + MatrixInverse, - Space: FunctionSpaceTrait, + Space: LocalFunctionSpaceTrait, Integrand: BoundaryIntegrand, K: Kernel, >( diff --git a/src/evaluator_tools.rs b/src/evaluator_tools.rs index 75c9ddd1..2bd02ffb 100644 --- a/src/evaluator_tools.rs +++ b/src/evaluator_tools.rs @@ -1,16 +1,24 @@ //! Various helper functions to support evaluators. -use itertools::izip; +use std::marker::PhantomData; + +use green_kernels::traits::Kernel; +use itertools::{izip, Itertools}; use mpi::traits::{Communicator, Equivalence}; use ndelement::traits::FiniteElement; use ndgrid::{ - traits::{Entity, GeometryMap, Grid}, + traits::{Entity, GeometryMap, Grid, ParallelGrid}, types::Ownership, }; + +use rayon::prelude::*; use rlst::{ + operator::interface::{ + distributed_sparse_operator::DistributedCsrMatrixOperator, DistributedArrayVectorSpace, + }, rlst_array_from_slice2, rlst_dynamic_array1, rlst_dynamic_array2, rlst_dynamic_array3, - rlst_dynamic_array4, DefaultIterator, DistributedCsrMatrix, IndexLayout, RawAccess, - RawAccessMut, RlstScalar, + rlst_dynamic_array4, AsApply, DefaultIterator, DistributedCsrMatrix, Element, IndexLayout, + OperatorBase, RawAccess, RawAccessMut, RlstScalar, }; use crate::function::FunctionSpaceTrait; @@ -18,19 +26,19 @@ use crate::function::FunctionSpaceTrait; /// Create a linear operator from the map of a basis to points. pub fn basis_to_point_map< 'a, - C: Communicator + 'a, + C: Communicator, DomainLayout: IndexLayout, RangeLayout: IndexLayout, T: RlstScalar + Equivalence, Space: FunctionSpaceTrait, >( function_space: &Space, - domain_layout: &'a DomainLayout, - range_layout: &'a RangeLayout, + domain_space: &'a DistributedArrayVectorSpace<'a, DomainLayout, T>, + range_space: &'a DistributedArrayVectorSpace<'a, RangeLayout, T>, quadrature_points: &[T::Real], quadrature_weights: &[T::Real], return_transpose: bool, -) -> DistributedCsrMatrix<'a, DomainLayout, RangeLayout, T, C> +) -> DistributedCsrMatrixOperator<'a, DomainLayout, RangeLayout, T, C> where T::Real: Equivalence, { @@ -62,8 +70,8 @@ where // Check that domain space and function space are compatible. - let n_domain_dofs = domain_layout.number_of_global_indices(); - let n_range_dofs = range_layout.number_of_global_indices(); + let n_domain_dofs = domain_space.index_layout().number_of_global_indices(); + let n_range_dofs = range_space.index_layout().number_of_global_indices(); assert_eq!(function_space.local_size(), n_domain_dofs,); @@ -128,31 +136,219 @@ where for (qindex, (jdet, qweight)) in izip!(jdets.iter(), quadrature_weights.iter()).enumerate() { - for global_dof in global_dofs.iter() { + for (i, global_dof) in global_dofs.iter().enumerate() { rows.push(n_quadrature_points * cell.global_index() + qindex); cols.push(*global_dof); - data.push(T::from_real(jdet * *qweight)); + data.push(T::from_real(jdet * *qweight) * table[[0, qindex, i, 0]]); } } } if return_transpose { - DistributedCsrMatrix::from_aij( - &domain_layout, - &range_layout, - &cols, - &rows, - &data, - domain_layout.comm(), + DistributedCsrMatrixOperator::new( + DistributedCsrMatrix::from_aij( + domain_space.index_layout(), + range_space.index_layout(), + &cols, + &rows, + &data, + ), + domain_space, + range_space, ) } else { - DistributedCsrMatrix::from_aij( - &domain_layout, - &range_layout, - &rows, - &cols, - &data, - domain_layout.comm(), + DistributedCsrMatrixOperator::new( + DistributedCsrMatrix::from_aij( + domain_space.index_layout(), + range_space.index_layout(), + &rows, + &cols, + &data, + ), + domain_space, + range_space, ) } } + +/// A linear operator that evaluates kernel interactions for a nonsingular quadrature rule on neighbouring elements. +/// +/// This creates a linear operator which for a given kernel evaluates all the interactions between neighbouring elements on a set +/// of local points per triangle. +pub struct NeighbourEvaluator< + 'a, + T: RlstScalar + Equivalence, + K: Kernel, + DomainLayout: IndexLayout, + RangeLayout: IndexLayout, + GridImpl: ParallelGrid, + C: Communicator, +> { + eval_points: Vec, + kernel: K, + domain_space: &'a DistributedArrayVectorSpace<'a, DomainLayout, T>, + range_space: &'a DistributedArrayVectorSpace<'a, RangeLayout, T>, + grid: &'a GridImpl, + active_cells: Vec, + _marker: PhantomData, +} + +impl< + 'a, + T: RlstScalar + Equivalence, + K: Kernel, + DomainLayout: IndexLayout, + RangeLayout: IndexLayout, + GridImpl: ParallelGrid, + C: Communicator, + > NeighbourEvaluator<'a, T, K, DomainLayout, RangeLayout, GridImpl, C> +{ + /// Create a new neighbour evaluator. + pub fn new( + eval_points: &[T::Real], + kernel: K, + domain_space: &'a DistributedArrayVectorSpace<'a, DomainLayout, T>, + range_space: &'a DistributedArrayVectorSpace<'a, RangeLayout, T>, + grid: &'a GridImpl, + ) -> Self { + // Check that the domain space and range space are compatible with the grid. + // Topological dimension of the grid. + let tdim = grid.topology_dim(); + + // The method is currently restricted to single element type grids. + // So let's make sure that this is the case. + + assert_eq!(grid.entity_types(tdim).len(), 1); + + // Get the number of points + assert_eq!(eval_points.len() % tdim, 0); + let n_points = eval_points.len() / tdim; + + // The active cells are those that we need to iterate over. + // At the moment these are simply all owned cells in the grid. + // We sort the active cells by global index. This is important so that in the evaluation + // we can just iterate the output vector through in chunks and know from the chunk which + // active cell it is associated with. + + let active_cells: Vec = grid + .entity_iter(tdim) + .filter(|e| matches!(e.ownership(), Ownership::Owned)) + .sorted_by_key(|e| e.global_index()) + .map(|e| e.local_index()) + .collect_vec(); + + let n_cells = active_cells.len(); + + // Check that domain space and function space are compatible with the grid. + + assert_eq!( + domain_space.index_layout().number_of_local_indices(), + n_cells * n_points + ); + + assert_eq!( + range_space.index_layout().number_of_local_indices(), + n_cells * n_points + ); + + Self { + eval_points: eval_points.to_vec(), + kernel, + domain_space, + range_space, + grid, + active_cells, + _marker: PhantomData, + } + } +} + +impl< + T: RlstScalar + Equivalence, + K: Kernel, + DomainLayout: IndexLayout, + RangeLayout: IndexLayout, + GridImpl: ParallelGrid, + C: Communicator, + > std::fmt::Debug for NeighbourEvaluator<'_, T, K, DomainLayout, RangeLayout, GridImpl, C> +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "Neighbourhood Evaluator with dimenion [{}, {}].", + self.range_space.index_layout().number_of_global_indices(), + self.domain_space.index_layout().number_of_global_indices() + ) + } +} + +impl< + 'a, + T: RlstScalar + Equivalence, + K: Kernel, + DomainLayout: IndexLayout, + RangeLayout: IndexLayout, + GridImpl: ParallelGrid, + C: Communicator, + > OperatorBase for NeighbourEvaluator<'a, T, K, DomainLayout, RangeLayout, GridImpl, C> +{ + type Domain = DistributedArrayVectorSpace<'a, DomainLayout, T>; + + type Range = DistributedArrayVectorSpace<'a, RangeLayout, T>; + + fn domain(&self) -> &Self::Domain { + self.domain_space + } + + fn range(&self) -> &Self::Range { + self.range_space + } +} + +impl< + T: RlstScalar + Equivalence, + K: Kernel, + DomainLayout: IndexLayout + Sync, + RangeLayout: IndexLayout + Sync, + GridImpl: ParallelGrid + Sync, + C: Communicator, + > AsApply for NeighbourEvaluator<'_, T, K, DomainLayout, RangeLayout, GridImpl, C> +{ + fn apply_extended( + &self, + alpha: ::F, + x: &::E, + beta: ::F, + y: &mut ::E, + ) -> rlst::RlstResult<()> { + // We need to iterate through the elements. + + let tdim = self.grid.topology_dim(); + + // In the new function we already made sure that eval_points is a multiple of tdim. + let n_points = self.eval_points.len() / tdim; + + // We go through groups of target dofs in chunks of lenth n_points. + // This corresponds to iteration in active cells since we ordered those + // already by global index. + + let raw_ptr = y.view_mut().local_mut().data_mut().as_mut_ptr(); + + // y.view_mut() + // .local_mut() + // .data_mut() + // .par_chunks_mut(n_points) + // .enumerate() + // .for_each(|(chunk_index, chunk)| { + // let cell = self.active_cells[chunk_index]; + // let cell_entity = self.grid.entity(tdim, cell).unwrap(); + + // // Get the geometry map for the cell. + // let geometry_map = self + // .grid + // .geometry_map(cell_entity.entity_type(), &self.eval_points); + // }); + + Ok(()) + } +} diff --git a/src/function.rs b/src/function.rs index c725505a..0f9bf0e2 100644 --- a/src/function.rs +++ b/src/function.rs @@ -17,30 +17,21 @@ use std::marker::PhantomData; type DofList = Vec>; type OwnerData = Vec<(usize, usize, usize, usize)>; -/// A function space -pub trait FunctionSpaceTrait { - /// Communicator - type C: Communicator; +/// A local function space +pub trait LocalFunctionSpaceTrait { /// Scalar type type T: RlstScalar; /// The grid type - type Grid: Grid::Real, EntityDescriptor = ReferenceCellType> - + ParallelGrid; + type LocalGrid: Grid::Real, EntityDescriptor = ReferenceCellType>; /// The finite element type type FiniteElement: FiniteElement + Sync; - /// Get the communicator - fn comm(&self) -> &Self::C; - /// Get the grid that the element is defined on - fn grid(&self) -> &Self::Grid; + fn grid(&self) -> &Self::LocalGrid; /// Get the finite element used to define this function space fn element(&self, cell_type: ReferenceCellType) -> &Self::FiniteElement; - /// Check if the function space is stored in serial - fn is_serial(&self) -> bool; - /// Get the DOF numbers on the local process associated with the given entity fn get_local_dof_numbers(&self, entity_dim: usize, entity_number: usize) -> &[usize]; @@ -69,12 +60,32 @@ pub trait FunctionSpaceTrait { fn ownership(&self, local_dof_index: usize) -> Ownership; } -/// Implementation of a general function space. -pub struct FunctionSpace< +/// A function space +pub trait FunctionSpaceTrait: LocalFunctionSpaceTrait { + /// Communicator + type C: Communicator; + + /// The grid type + type Grid: ParallelGrid; + /// Local Function Space + type LocalFunctionSpace: LocalFunctionSpaceTrait< + T = Self::T, + LocalGrid = Self::LocalGrid, + FiniteElement = Self::FiniteElement, + >; + + /// Get the communicator + fn comm(&self) -> &Self::C; + + /// Get the local function space + fn local_space(&self) -> &Self::LocalFunctionSpace; +} + +/// Definition of a local function space. +pub struct LocalFunctionSpace< 'a, - C: Communicator, T: RlstScalar + MatrixInverse, - GridImpl: ParallelGrid + Grid, + GridImpl: Grid, > { grid: &'a GridImpl, elements: HashMap>, @@ -84,23 +95,178 @@ pub struct FunctionSpace< global_size: usize, global_dof_numbers: Vec, ownership: Vec, - _marker: std::marker::PhantomData, } -unsafe impl< - C: Communicator, +impl< + 'a, T: RlstScalar + MatrixInverse, - GridImpl: ParallelGrid + Grid, - > Sync for FunctionSpace<'_, C, T, GridImpl> + GridImpl: Grid, + > LocalFunctionSpace<'a, T, GridImpl> { + /// Create new local function space + #[allow(clippy::too_many_arguments)] + pub fn new( + grid: &'a GridImpl, + elements: HashMap>, + entity_dofs: [Vec>; 4], + cell_dofs: Vec>, + local_size: usize, + global_size: usize, + global_dof_numbers: Vec, + ownership: Vec, + ) -> Self { + Self { + grid, + elements, + entity_dofs, + cell_dofs, + local_size, + global_size, + global_dof_numbers, + ownership, + } + } +} + +impl< + T: RlstScalar + MatrixInverse, + GridImpl: Grid, + > LocalFunctionSpaceTrait for LocalFunctionSpace<'_, T, GridImpl> +{ + type T = T; + + type LocalGrid = GridImpl; + + type FiniteElement = CiarletElement; + + fn grid(&self) -> &Self::LocalGrid { + self.grid + } + + fn element(&self, cell_type: ReferenceCellType) -> &Self::FiniteElement { + &self.elements[&cell_type] + } + + fn get_local_dof_numbers(&self, entity_dim: usize, entity_number: usize) -> &[usize] { + &self.entity_dofs[entity_dim][entity_number] + } + + fn local_size(&self) -> usize { + self.local_size + } + + fn global_size(&self) -> usize { + self.global_size + } + unsafe fn cell_dofs_unchecked(&self, cell: usize) -> &[usize] { + self.cell_dofs.get_unchecked(cell) + } + fn cell_dofs(&self, cell: usize) -> Option<&[usize]> { + if cell < self.cell_dofs.len() { + Some(unsafe { self.cell_dofs_unchecked(cell) }) + } else { + None + } + } + fn cell_colouring(&self) -> HashMap>> { + let mut colouring = HashMap::new(); + //: HashMap>> + for cell in self.grid.entity_types(2) { + colouring.insert(*cell, vec![]); + } + let mut edim = 0; + while self.elements[&self.grid.entity_types(2)[0]] + .entity_dofs(edim, 0) + .unwrap() + .is_empty() + { + edim += 1; + } + + let mut entity_colours = vec![ + vec![]; + if edim == 0 { + self.grid.entity_count(ReferenceCellType::Point) + } else if edim == 1 { + self.grid.entity_count(ReferenceCellType::Interval) + } else if edim == 2 && self.grid.topology_dim() == 2 { + self.grid + .entity_types(2) + .iter() + .map(|&i| self.grid.entity_count(i)) + .sum::() + } else { + unimplemented!(); + } + ]; + + for cell in self.grid.entity_iter(2) { + let cell_type = cell.entity_type(); + let indices = cell.topology().sub_entity_iter(edim).collect::>(); + + let c = { + let mut c = 0; + while c < colouring[&cell_type].len() { + let mut found = false; + for v in &indices { + if entity_colours[*v].contains(&c) { + found = true; + break; + } + } + + if !found { + break; + } + c += 1; + } + c + }; + if c == colouring[&cell_type].len() { + for ct in self.grid.entity_types(2) { + colouring.get_mut(ct).unwrap().push(if *ct == cell_type { + vec![cell.local_index()] + } else { + vec![] + }); + } + } else { + colouring.get_mut(&cell_type).unwrap()[c].push(cell.local_index()); + } + for v in &indices { + entity_colours[*v].push(c); + } + } + colouring + } + fn global_dof_index(&self, local_dof_index: usize) -> usize { + self.global_dof_numbers[local_dof_index] + } + fn ownership(&self, local_dof_index: usize) -> Ownership { + self.ownership[local_dof_index] + } +} + +/// Implementation of a general function space. +pub struct FunctionSpace< + 'a, + C: Communicator, + T: RlstScalar + MatrixInverse, + GridImpl: ParallelGrid, +> { + grid: &'a GridImpl, + local_space: LocalFunctionSpace<'a, T, GridImpl::LocalGrid>, + _marker: PhantomData<(C, T)>, } impl< 'a, C: Communicator, T: RlstScalar + MatrixInverse, - GridImpl: ParallelGrid + Grid, + GridImpl: ParallelGrid, > FunctionSpace<'a, C, T, GridImpl> +where + GridImpl::LocalGrid: Sync, { /// Create new function space pub fn new( @@ -225,140 +391,108 @@ impl< Self { grid, - elements, - entity_dofs, - cell_dofs, - local_size: dofmap_size, - global_size, - global_dof_numbers, - ownership, + local_space: LocalFunctionSpace::new( + grid.local_grid(), + elements, + entity_dofs, + cell_dofs, + dofmap_size, + global_size, + global_dof_numbers, + ownership, + ), _marker: PhantomData, } + + // Self { + // grid, + // elements, + // entity_dofs, + // cell_dofs, + // local_size: dofmap_size, + // global_size, + // global_dof_numbers, + // ownership, + // _marker: PhantomData, + // } } } impl< C: Communicator, T: RlstScalar + MatrixInverse, - GridImpl: ParallelGrid + Grid, - > FunctionSpaceTrait for FunctionSpace<'_, C, T, GridImpl> + GridImpl: ParallelGrid, + > LocalFunctionSpaceTrait for FunctionSpace<'_, C, T, GridImpl> { type T = T; - type Grid = GridImpl; + + type LocalGrid = GridImpl::LocalGrid; + type FiniteElement = CiarletElement; - type C = C; - fn grid(&self) -> &Self::Grid { - self.grid + fn grid(&self) -> &Self::LocalGrid { + self.local_space.grid() } - fn element(&self, cell_type: ReferenceCellType) -> &CiarletElement { - &self.elements[&cell_type] + + fn element(&self, cell_type: ReferenceCellType) -> &Self::FiniteElement { + self.local_space.element(cell_type) } fn get_local_dof_numbers(&self, entity_dim: usize, entity_number: usize) -> &[usize] { - &self.entity_dofs[entity_dim][entity_number] - } - fn is_serial(&self) -> bool { - self.grid.comm().size() == 1 + self.local_space + .get_local_dof_numbers(entity_dim, entity_number) } + fn local_size(&self) -> usize { - self.local_size + self.local_space.local_size() } fn global_size(&self) -> usize { - self.global_size - } - unsafe fn cell_dofs_unchecked(&self, cell: usize) -> &[usize] { - self.cell_dofs.get_unchecked(cell) + self.local_space.global_size() } + fn cell_dofs(&self, cell: usize) -> Option<&[usize]> { - if cell < self.cell_dofs.len() { - Some(unsafe { self.cell_dofs_unchecked(cell) }) - } else { - None - } + self.local_space.cell_dofs(cell) } - fn cell_colouring(&self) -> HashMap>> { - let mut colouring = HashMap::new(); - //: HashMap>> - for cell in self.grid.entity_types(2) { - colouring.insert(*cell, vec![]); - } - let mut edim = 0; - while self.elements[&self.grid.entity_types(2)[0]] - .entity_dofs(edim, 0) - .unwrap() - .is_empty() - { - edim += 1; - } - - let mut entity_colours = vec![ - vec![]; - if edim == 0 { - self.grid.entity_count(ReferenceCellType::Point) - } else if edim == 1 { - self.grid.entity_count(ReferenceCellType::Interval) - } else if edim == 2 && self.grid.topology_dim() == 2 { - self.grid - .entity_types(2) - .iter() - .map(|&i| self.grid.entity_count(i)) - .sum::() - } else { - unimplemented!(); - } - ]; - for cell in self.grid.entity_iter(2) { - let cell_type = cell.entity_type(); - let indices = cell.topology().sub_entity_iter(edim).collect::>(); - - let c = { - let mut c = 0; - while c < colouring[&cell_type].len() { - let mut found = false; - for v in &indices { - if entity_colours[*v].contains(&c) { - found = true; - break; - } - } + unsafe fn cell_dofs_unchecked(&self, cell: usize) -> &[usize] { + self.local_space.cell_dofs_unchecked(cell) + } - if !found { - break; - } - c += 1; - } - c - }; - if c == colouring[&cell_type].len() { - for ct in self.grid.entity_types(2) { - colouring.get_mut(ct).unwrap().push(if *ct == cell_type { - vec![cell.local_index()] - } else { - vec![] - }); - } - } else { - colouring.get_mut(&cell_type).unwrap()[c].push(cell.local_index()); - } - for v in &indices { - entity_colours[*v].push(c); - } - } - colouring + fn cell_colouring(&self) -> HashMap>> { + self.local_space.cell_colouring() } + fn global_dof_index(&self, local_dof_index: usize) -> usize { - self.global_dof_numbers[local_dof_index] + self.local_space.global_dof_index(local_dof_index) } + fn ownership(&self, local_dof_index: usize) -> Ownership { - self.ownership[local_dof_index] + // Syntactical workaround as rust-analyzer mixed up this ownership with entity ownership. + LocalFunctionSpaceTrait::ownership(&self.local_space, local_dof_index) } +} +impl< + 'a, + C: Communicator, + T: RlstScalar + MatrixInverse, + GridImpl: ParallelGrid, + > FunctionSpaceTrait for FunctionSpace<'a, C, T, GridImpl> +{ fn comm(&self) -> &C { self.grid.comm() } + + type C = C; + + type Grid = GridImpl; + + type LocalFunctionSpace = LocalFunctionSpace<'a, T, GridImpl::LocalGrid>; + + fn local_space(&self) -> &Self::LocalFunctionSpace { + &self.local_space + } } /// Assign DOFs to entities. diff --git a/src/greens_function_evaluators/dense_evaluator.rs b/src/greens_function_evaluators/dense_evaluator.rs index 8c70faac..6260d479 100644 --- a/src/greens_function_evaluators/dense_evaluator.rs +++ b/src/greens_function_evaluators/dense_evaluator.rs @@ -28,13 +28,12 @@ pub struct DenseEvaluator< } impl< - 'a, C: Communicator, T: RlstScalar + Equivalence, SourceLayout: IndexLayout, TargetLayout: IndexLayout, K: DistributedKernelEvaluator, - > std::fmt::Debug for DenseEvaluator<'a, C, T, SourceLayout, TargetLayout, K> + > std::fmt::Debug for DenseEvaluator<'_, C, T, SourceLayout, TargetLayout, K> where T::Real: Equivalence, { @@ -138,13 +137,12 @@ where } impl< - 'a, C: Communicator, T: RlstScalar + Equivalence, SourceLayout: IndexLayout, TargetLayout: IndexLayout, K: DistributedKernelEvaluator, - > AsApply for DenseEvaluator<'a, C, T, SourceLayout, TargetLayout, K> + > AsApply for DenseEvaluator<'_, C, T, SourceLayout, TargetLayout, K> where T::Real: Equivalence, { diff --git a/src/greens_function_evaluators/kifmm_evaluator.rs b/src/greens_function_evaluators/kifmm_evaluator.rs index 6abbacec..86720cee 100644 --- a/src/greens_function_evaluators/kifmm_evaluator.rs +++ b/src/greens_function_evaluators/kifmm_evaluator.rs @@ -51,12 +51,11 @@ pub struct KiFmmEvaluator< } impl< - 'a, C: Communicator, T: RlstScalar + Equivalence, SourceLayout: IndexLayout, TargetLayout: IndexLayout, - > std::fmt::Debug for KiFmmEvaluator<'a, C, T, SourceLayout, TargetLayout> + > std::fmt::Debug for KiFmmEvaluator<'_, C, T, SourceLayout, TargetLayout> where T::Real: Equivalence, T: DftType::ComplexType>, @@ -197,12 +196,11 @@ where } impl< - 'a, C: Communicator, T: RlstScalar + Equivalence, SourceLayout: IndexLayout, TargetLayout: IndexLayout, - > AsApply for KiFmmEvaluator<'a, C, T, SourceLayout, TargetLayout> + > AsApply for KiFmmEvaluator<'_, C, T, SourceLayout, TargetLayout> where T::Real: Equivalence, T: DftType::ComplexType>, diff --git a/tests/test_space.rs b/tests/test_space.rs index 6d1d9e78..d736de63 100644 --- a/tests/test_space.rs +++ b/tests/test_space.rs @@ -1,8 +1,9 @@ -use bempp::function::{assign_dofs, FunctionSpace, FunctionSpaceTrait}; +use bempp::function::{assign_dofs, FunctionSpace, LocalFunctionSpaceTrait}; use bempp::shapes::{regular_sphere, screen_triangles}; +use mpi::traits::Communicator; use ndelement::ciarlet::{LagrangeElementFamily, RaviartThomasElementFamily}; use ndelement::types::{Continuity, ReferenceCellType}; -use ndgrid::traits::{Entity, Grid, Topology}; +use ndgrid::traits::{Entity, Grid, ParallelGrid, Topology}; use std::sync::LazyLock; use mpi::environment::Universe; @@ -13,13 +14,18 @@ static MPI_UNIVERSE: LazyLock = std::sync::LazyLock::new(|| { .0 }); -fn run_test( - grid: &(impl Grid + Sync), +fn run_test< + C: Communicator, + GridImpl: ParallelGrid, +>( + grid: &GridImpl, degree: usize, continuity: Continuity, -) { +) where + GridImpl::LocalGrid: Sync, +{ let family = LagrangeElementFamily::::new(degree, continuity); - let (cell_dofs, entity_dofs, size, owner_data) = assign_dofs(0, grid, &family); + let (cell_dofs, entity_dofs, size, owner_data) = assign_dofs(0, grid.local_grid(), &family); for o in &owner_data { assert_eq!(o.0, 0); @@ -41,13 +47,18 @@ fn run_test( } } -fn run_test_rt( - grid: &(impl Grid + Sync), +fn run_test_rt< + C: Communicator, + GridImpl: ParallelGrid, +>( + grid: &GridImpl, degree: usize, continuity: Continuity, -) { +) where + GridImpl::LocalGrid: Sync, +{ let family = RaviartThomasElementFamily::::new(degree, continuity); - let (cell_dofs, entity_dofs, size, owner_data) = assign_dofs(0, grid, &family); + let (cell_dofs, entity_dofs, size, owner_data) = assign_dofs(0, grid.local_grid(), &family); for o in &owner_data { assert_eq!(o.0, 0);