Parallel function space does not depend on serial space anymore.

src/function/function_space/parallel.rs

@@ -1,18 +1,16 @@
 //! Parallel function space
 
-use crate::function::{
-    function_space::assign_dofs, FunctionSpace, MPIFunctionSpace, SerialFunctionSpace,
-};
+use crate::function::{function_space::assign_dofs, FunctionSpace, MPIFunctionSpace};
 use mpi::{
     point_to_point::{Destination, Source},
     request::WaitGuard,
     topology::Communicator,
 };
-use ndelement::ciarlet::CiarletElement;
 use ndelement::traits::ElementFamily;
 use ndelement::types::ReferenceCellType;
+use ndelement::{ciarlet::CiarletElement, traits::FiniteElement};
 use ndgrid::{
-    traits::{Grid, ParallelGrid},
+    traits::{Entity, Grid, ParallelGrid, Topology},
     types::Ownership,
 };
 use rlst::{MatrixInverse, RlstScalar};
@@ -24,12 +22,54 @@ pub struct LocalFunctionSpace<
     T: RlstScalar + MatrixInverse,
     GridImpl: Grid<T = T::Real, EntityDescriptor = ReferenceCellType> + Sync,
 > {
-    serial_space: SerialFunctionSpace<'a, T, GridImpl>,
+    grid: &'a GridImpl,
+    elements: HashMap<ReferenceCellType, CiarletElement<T>>,
+    entity_dofs: [Vec<Vec<usize>>; 4],
+    cell_dofs: Vec<Vec<usize>>,
+    local_size: usize,
     global_size: usize,
     global_dof_numbers: Vec<usize>,
     ownership: Vec<Ownership>,
 }
 
+impl<
+        'a,
+        T: RlstScalar + MatrixInverse,
+        GridImpl: Grid<T = T::Real, EntityDescriptor = ReferenceCellType> + Sync,
+    > LocalFunctionSpace<'a, T, GridImpl>
+{
+    // /// Create new function space
+    // pub fn new(
+    //     grid: &'a GridImpl,
+    //     e_family: &impl ElementFamily<
+    //         T = T,
+    //         FiniteElement = CiarletElement<T>,
+    //         CellType = ReferenceCellType,
+    //     >,
+    //     global_size: usize,
+    //     global_dof_numbers: Vec<usize>,
+    //     ownership: Vec<Ownership>,
+    // ) -> Self {
+    //     let (cell_dofs, entity_dofs, local_size, _) = assign_dofs(0, grid, e_family);
+
+    //     let mut elements = HashMap::new();
+    //     for cell in grid.entity_types(2) {
+    //         elements.insert(*cell, e_family.element(*cell));
+    //     }
+
+    //     Self {
+    //         grid,
+    //         elements,
+    //         entity_dofs,
+    //         cell_dofs,
+    //         local_size,
+    //         global_size,
+    //         global_dof_numbers,
+    //         ownership,
+    //     }
+    // }
+}
+
 impl<
         T: RlstScalar + MatrixInverse,
         GridImpl: Grid<T = T::Real, EntityDescriptor = ReferenceCellType> + Sync,
@@ -47,33 +87,105 @@ impl<
         self
     }
     fn grid(&self) -> &Self::Grid {
-        self.serial_space.grid
+        self.grid
     }
     fn element(&self, cell_type: ReferenceCellType) -> &CiarletElement<T> {
-        self.serial_space.element(cell_type)
+        &self.elements[&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)
+        &self.entity_dofs[entity_dim][entity_number]
     }
     fn is_serial(&self) -> bool {
         false
     }
     fn local_size(&self) -> usize {
-        self.serial_space.local_size()
+        self.local_size
     }
     fn global_size(&self) -> usize {
         self.global_size
     }
     unsafe fn cell_dofs_unchecked(&self, cell: usize) -> &[usize] {
-        self.serial_space.cell_dofs_unchecked(cell)
+        self.cell_dofs.get_unchecked(cell)
     }
     fn cell_dofs(&self, cell: usize) -> Option<&[usize]> {
-        self.serial_space.cell_dofs(cell)
+        if cell < self.cell_dofs.len() {
+            Some(unsafe { self.cell_dofs_unchecked(cell) })
+        } else {
+            None
+        }
     }
     fn cell_colouring(&self) -> HashMap<ReferenceCellType, Vec<Vec<usize>>> {
-        self.serial_space.cell_colouring()
+        let mut colouring = HashMap::new();
+        //: HashMap<ReferenceCellType, Vec<Vec<usize>>>
+        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::<usize>()
+            } 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::<Vec<_>>();
+
+            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]
     }
@@ -220,15 +332,12 @@ impl<
             }
         }
 
-        let serial_space = SerialFunctionSpace {
+        let local_space = LocalFunctionSpace {
             grid: grid.local_grid(),
             elements,
             entity_dofs,
             cell_dofs,
-            size: dofmap_size,
-        };
-        let local_space = LocalFunctionSpace {
-            serial_space,
+            local_size: dofmap_size,
             global_size,
             global_dof_numbers,
             ownership,