diff --git a/src/assembly/batched.rs b/src/assembly/batched.rs index 025ea62a..c98914eb 100644 --- a/src/assembly/batched.rs +++ b/src/assembly/batched.rs @@ -12,7 +12,6 @@ use crate::traits::function::FunctionSpace; #[cfg(feature = "mpi")] use crate::traits::function::FunctionSpaceInParallel; use crate::traits::grid::{CellType, GridType, ReferenceMapType, TopologyType}; -#[cfg(feature = "mpi")] use crate::traits::types::Ownership; use crate::traits::types::ReferenceCellType; use green_kernels::types::EvalType; @@ -211,9 +210,11 @@ fn assemble_batch_singular< .unwrap(); } } - output.rows.push(*test_dof); - output.cols.push(*trial_dof); - output.data.push(sum); + if trial_space.ownership(*trial_dof) == Ownership::Owned { + output.rows.push(test_space.global_dof_index(*test_dof)); + output.cols.push(trial_space.global_dof_index(*trial_dof)); + output.data.push(sum); + } } } } @@ -471,9 +472,11 @@ fn assemble_batch_singular_correction< }; } } - output.rows.push(*test_dof); - output.cols.push(*trial_dof); - output.data.push(sum); + if trial_space.ownership(*trial_dof) == Ownership::Owned { + output.rows.push(test_space.global_dof_index(*test_dof)); + output.cols.push(trial_space.global_dof_index(*trial_dof)); + output.data.push(sum); + } } } } @@ -1042,34 +1045,7 @@ pub trait BatchedAssembler: Sync + Sized { trial_space: &(impl FunctionSpaceInParallel + FunctionSpace), test_space: &(impl FunctionSpaceInParallel + FunctionSpace), ) -> CsrMatrix { - let local_matrix = - self.assemble_singular_into_csr(trial_space.local_space(), test_space.local_space()); - - let ownership = trial_space.ownership(); - let test_global_dof_numbers = test_space.global_dof_numbers(); - let trial_global_dof_numbers = trial_space.global_dof_numbers(); - - let mut rows = vec![]; - let mut cols = vec![]; - let mut data = vec![]; - let mut r = 0; - for (i, index) in local_matrix.indices().iter().enumerate() { - while i >= local_matrix.indptr()[r + 1] { - r += 1; - } - if ownership[*index] == Ownership::Owned { - rows.push(test_global_dof_numbers[r]); - cols.push(trial_global_dof_numbers[*index]); - data.push(local_matrix.data()[i]); - } - } - CsrMatrix::from_aij( - [test_space.global_size(), trial_space.global_size()], - &rows, - &cols, - &data, - ) - .unwrap() + self.assemble_singular_into_csr(trial_space.local_space(), test_space.local_space()) } /// Assemble the singular correction into a dense matrix @@ -1139,36 +1115,10 @@ pub trait BatchedAssembler: Sync + Sized { trial_space: &(impl FunctionSpaceInParallel + FunctionSpace), test_space: &(impl FunctionSpaceInParallel + FunctionSpace), ) -> CsrMatrix { - let local_matrix = self.assemble_singular_correction_into_csr( + self.assemble_singular_correction_into_csr( trial_space.local_space(), test_space.local_space(), - ); - - let ownership = trial_space.ownership(); - let test_global_dof_numbers = test_space.global_dof_numbers(); - let trial_global_dof_numbers = trial_space.global_dof_numbers(); - - let mut rows = vec![]; - let mut cols = vec![]; - let mut data = vec![]; - let mut r = 0; - for (i, index) in local_matrix.indices().iter().enumerate() { - while i >= local_matrix.indptr()[r + 1] { - r += 1; - } - if ownership[*index] == Ownership::Owned { - rows.push(test_global_dof_numbers[r]); - cols.push(trial_global_dof_numbers[*index]); - data.push(local_matrix.data()[i]); - } - } - CsrMatrix::from_aij( - [test_space.global_size(), trial_space.global_size()], - &rows, - &cols, - &data, ) - .unwrap() } /// Assemble into a dense matrix diff --git a/src/function/function_space/parallel.rs b/src/function/function_space/parallel.rs index 428e27e3..62ef6cee 100644 --- a/src/function/function_space/parallel.rs +++ b/src/function/function_space/parallel.rs @@ -16,6 +16,50 @@ use mpi::{ use rlst::RlstScalar; use std::collections::HashMap; +/// The local function space on a process +pub struct LocalFunctionSpace<'a, T: RlstScalar, GridImpl: GridType> { + serial_space: SerialFunctionSpace<'a, T, GridImpl>, + global_size: usize, + global_dof_numbers: Vec, + ownership: Vec, +} + +impl<'a, T: RlstScalar, GridImpl: GridType> FunctionSpace + for LocalFunctionSpace<'a, T, GridImpl> +{ + type Grid = GridImpl; + type FiniteElement = CiarletElement; + + fn grid(&self) -> &Self::Grid { + self.serial_space.grid + } + fn element(&self, cell_type: ReferenceCellType) -> &CiarletElement { + self.serial_space.element(cell_type) + } + fn get_local_dof_numbers(&self, entity_dim: usize, entity_number: usize) -> &[usize] { + self.serial_space + .get_local_dof_numbers(entity_dim, entity_number) + } + fn local_size(&self) -> usize { + self.serial_space.local_size() + } + fn global_size(&self) -> usize { + self.global_size + } + fn cell_dofs(&self, cell: usize) -> Option<&[usize]> { + self.serial_space.cell_dofs(cell) + } + fn cell_colouring(&self) -> HashMap>> { + self.serial_space.cell_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] + } +} + /// A parallel function space pub struct ParallelFunctionSpace< 'a, @@ -23,10 +67,7 @@ pub struct ParallelFunctionSpace< GridImpl: ParallelGridType + GridType, > { grid: &'a GridImpl, - local_space: SerialFunctionSpace<'a, T, ::LocalGridType>, - global_dof_numbers: Vec, - owners: Vec, - global_size: usize, + local_space: LocalFunctionSpace<'a, T, ::LocalGridType>, } impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> @@ -50,16 +91,8 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> elements.insert(*cell, e_family.element(*cell)); } - let local_space = SerialFunctionSpace { - grid: grid.local_grid(), - elements, - entity_dofs, - cell_dofs, - size: dofmap_size, - }; - // Assign global DOF numbers - let mut global_dof_numbers = vec![0; local_space.local_size()]; + let mut global_dof_numbers = vec![0; dofmap_size]; let mut ghost_indices = vec![vec![]; size as usize]; let mut ghost_cells = vec![vec![]; size as usize]; let mut ghost_cell_dofs = vec![vec![]; size as usize]; @@ -123,7 +156,7 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> let local_ghost_dofs = gcells .iter() .zip(&gdofs) - .map(|(c, d)| local_space.cell_dofs(*c).unwrap()[*d]) + .map(|(c, d)| cell_dofs[*c][*d]) .collect::>(); let global_ghost_dofs = local_ghost_dofs .iter() @@ -136,7 +169,7 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> } } // receive ghost info - let mut owners = vec![Ownership::Owned; local_space.local_size()]; + let mut ownership = vec![Ownership::Owned; dofmap_size]; for p in 0..size { if p != rank { let process = comm.process_at_rank(p); @@ -147,18 +180,26 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> .zip(local_ghost_dofs.iter().zip(&global_ghost_dofs)) { global_dof_numbers[*i] = *g; - owners[*i] = Ownership::Ghost(p as usize, *l); + ownership[*i] = Ownership::Ghost(p as usize, *l); } } } - Self { - grid, - local_space, - global_dof_numbers, - owners, + let serial_space = SerialFunctionSpace { + grid: grid.local_grid(), + elements, + entity_dofs, + cell_dofs, + size: dofmap_size, + }; + let local_space = LocalFunctionSpace { + serial_space, global_size, - } + global_dof_numbers, + ownership, + }; + + Self { grid, local_space } } } @@ -166,24 +207,14 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> Func for ParallelFunctionSpace<'a, T, GridImpl> { type ParallelGrid = GridImpl; - type SerialSpace = SerialFunctionSpace<'a, T, ::LocalGridType>; + type SerialSpace = LocalFunctionSpace<'a, T, ::LocalGridType>; /// Get the local space on the process fn local_space( &self, - ) -> &SerialFunctionSpace<'a, T, ::LocalGridType> { + ) -> &LocalFunctionSpace<'a, T, ::LocalGridType> { &self.local_space } - - /// Get the global DOF number associated with each local DOF - fn global_dof_numbers(&self) -> &Vec { - &self.global_dof_numbers - } - - /// Get ownership info - fn ownership(&self) -> &Vec { - &self.owners - } } impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> FunctionSpace @@ -206,7 +237,7 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> Func self.local_space.local_size() } fn global_size(&self) -> usize { - self.global_size + self.local_space.global_size() } fn cell_dofs(&self, cell: usize) -> Option<&[usize]> { self.local_space.cell_dofs(cell) @@ -214,4 +245,10 @@ impl<'a, T: RlstScalar, GridImpl: ParallelGridType + GridType> Func fn cell_colouring(&self) -> HashMap>> { self.local_space.cell_colouring() } + fn global_dof_index(&self, local_dof_index: usize) -> usize { + self.local_space.global_dof_index(local_dof_index) + } + fn ownership(&self, local_dof_index: usize) -> Ownership { + self.local_space.ownership(local_dof_index) + } } diff --git a/src/function/function_space/serial.rs b/src/function/function_space/serial.rs index 485ad387..808a37dc 100644 --- a/src/function/function_space/serial.rs +++ b/src/function/function_space/serial.rs @@ -6,7 +6,7 @@ use crate::traits::{ element::{ElementFamily, FiniteElement}, function::FunctionSpace, grid::{CellType, GridType, TopologyType}, - types::ReferenceCellType, + types::{Ownership, ReferenceCellType}, }; use rlst::RlstScalar; use std::collections::HashMap; @@ -62,7 +62,7 @@ impl<'a, T: RlstScalar, GridImpl: GridType> FunctionSpace self.size } fn global_size(&self) -> usize { - self.local_size() + self.size } fn cell_dofs(&self, cell: usize) -> Option<&[usize]> { if cell < self.cell_dofs.len() { @@ -146,6 +146,12 @@ impl<'a, T: RlstScalar, GridImpl: GridType> FunctionSpace } colouring } + fn global_dof_index(&self, local_dof_index: usize) -> usize { + local_dof_index + } + fn ownership(&self, _local_dof_index: usize) -> Ownership { + Ownership::Owned + } } #[cfg(test)] diff --git a/src/traits/function.rs b/src/traits/function.rs index ed248ed9..f4d901ca 100644 --- a/src/traits/function.rs +++ b/src/traits/function.rs @@ -3,7 +3,6 @@ use crate::traits::element::FiniteElement; use crate::traits::grid::GridType; #[cfg(feature = "mpi")] use crate::traits::grid::ParallelGridType; -#[cfg(feature = "mpi")] use crate::traits::types::Ownership; use crate::traits::types::ReferenceCellType; use std::collections::HashMap; @@ -40,6 +39,12 @@ pub trait FunctionSpace { /// Compute a colouring of the cells so that no two cells that share an entity with DOFs associated with it are assigned the same colour fn cell_colouring(&self) -> HashMap>>; + + /// Get the global DOF indes associated with a local DOF indec + fn global_dof_index(&self, local_dof_index: usize) -> usize; + + /// Get ownership of a local DOF + fn ownership(&self, local_dof_index: usize) -> Ownership; } #[cfg(feature = "mpi")] @@ -52,10 +57,4 @@ pub trait FunctionSpaceInParallel { /// Get the local space on the process fn local_space(&self) -> &Self::SerialSpace; - - /// Get the global DOF number associated with each local DOF - fn global_dof_numbers(&self) -> &Vec; - - /// Get ownership info - fn ownership(&self) -> &Vec; }