Repository: bempp/bempp-rs (language: python)

Issue:
Refactor boundary assembly into better traits
New trait structure (defined in src/traits/assembly.rs):
- `KernelEvaluator` evaluates Green's kernels with the correct EvalType
- `BoundaryIntegrand` evaluates the integrand for a given test and trial point on a given cell
- `CellPairAssembler` assembles the contribution to the matrix due to a pair of cells
- `BoundaryAssembler` and `ParallelBoundaryAssembler` contain user-facing full assembly functions
Convert BoundaryAssembler to struct
Follows on from #279.

Converts the trait BoundaryAssembler to a struct


Discussion / hints:

Fix patch:
diff --git a/src/assembly/potential.rs b/src/assembly/potential.rs
index 207c2117..ab2ff885 100644
--- a/src/assembly/potential.rs
+++ b/src/assembly/potential.rs
@@ -1,246 +1,6 @@
 //! Assembly of potential operators
-pub(crate) mod double_layer;
-pub(crate) mod single_layer;
+mod assemblers;
+pub(crate) mod cell_assemblers;
+mod integrands;
 
-pub use double_layer::DoubleLayerPotentialAssembler;
-pub use single_layer::SingleLayerPotentialAssembler;
-
-use crate::assembly::common::RawData2D;
-use crate::assembly::common::RlstArray;
-use crate::quadrature::simplex_rules::simplex_rule;
-use crate::traits::FunctionSpace;
-use ndelement::traits::FiniteElement;
-use ndelement::types::ReferenceCellType;
-use ndgrid::traits::{GeometryMap, Grid};
-use rayon::prelude::*;
-use rlst::{
-    rlst_dynamic_array2, rlst_dynamic_array3, rlst_dynamic_array4, MatrixInverse, RandomAccessMut,
-    RawAccess, RawAccessMut, RlstScalar, Shape, UnsafeRandomAccessByRef,
-};
-use std::collections::HashMap;
-
-/// Assemble the contribution to the terms of a matrix for a batch of non-adjacent cells
-#[allow(clippy::too_many_arguments)]
-fn assemble_batch<
-    T: RlstScalar + MatrixInverse,
-    G: Grid<T = T::Real, EntityDescriptor = ReferenceCellType>,
-    Element: FiniteElement<T = T> + Sync,
->(
-    assembler: &impl PotentialAssembler<T = T>,
-    deriv_size: usize,
-    output: &RawData2D<T>,
-    cell_type: ReferenceCellType,
-    space: &impl FunctionSpace<Grid = G, FiniteElement = Element>,
-    evaluation_points: &RlstArray<T::Real, 2>,
-    cells: &[usize],
-    points: &RlstArray<T::Real, 2>,
-    weights: &[T::Real],
-    table: &RlstArray<T, 4>,
-) -> usize {
-    let npts = weights.len();
-    let nevalpts = evaluation_points.shape()[1];
-    debug_assert!(points.shape()[1] == npts);
-
-    let grid = space.grid();
-
-    assert_eq!(grid.geometry_dim(), 3);
-    assert_eq!(grid.topology_dim(), 2);
-
-    let mut k = rlst_dynamic_array3!(T, [deriv_size, npts, nevalpts]);
-    let zero = num::cast::<f64, T::Real>(0.0).unwrap();
-    let mut jdet = vec![zero; npts];
-    let mut mapped_pts = rlst_dynamic_array2!(T::Real, [3, npts]);
-    let mut normals = rlst_dynamic_array2!(T::Real, [3, npts]);
-    let mut jacobians = rlst_dynamic_array2!(T::Real, [6, npts]);
-
-    let evaluator = grid.geometry_map(cell_type, points.data());
-
-    let mut sum: T;
-
-    for cell in cells {
-        evaluator.points(*cell, mapped_pts.data_mut());
-        evaluator.jacobians_dets_normals(
-            *cell,
-            jacobians.data_mut(),
-            &mut jdet,
-            normals.data_mut(),
-        );
-
-        assembler.kernel_assemble_st(mapped_pts.data(), evaluation_points.data(), k.data_mut());
-
-        let dofs = space.cell_dofs(*cell).unwrap();
-
-        for (i, dof) in dofs.iter().enumerate() {
-            for eval_index in 0..nevalpts {
-                sum = num::cast::<f64, T>(0.0).unwrap();
-                for (index, wt) in weights.iter().enumerate() {
-                    sum += unsafe {
-                        assembler.kernel_value(&k, &normals, index, eval_index)
-                            * num::cast::<T::Real, T>(*wt * jdet[index]).unwrap()
-                            * *table.get_unchecked([0, index, i, 0])
-                    };
-                }
-                unsafe {
-                    *output.data.add(eval_index + output.shape[0] * *dof) += sum;
-                }
-            }
-        }
-    }
-    1
-}
-
-/// Options for a potential assembler
-pub struct PotentialAssemblerOptions {
-    /// Number of points used in quadrature for non-singular integrals
-    quadrature_degrees: HashMap<ReferenceCellType, usize>,
-    /// Maximum size of each batch of cells to send to an assembly function
-    batch_size: usize,
-}
-
-impl Default for PotentialAssemblerOptions {
-    fn default() -> Self {
-        use ReferenceCellType::{Quadrilateral, Triangle};
-        Self {
-            quadrature_degrees: HashMap::from([(Triangle, 37), (Quadrilateral, 37)]),
-            batch_size: 128,
-        }
-    }
-}
-
-pub trait PotentialAssembler: Sync + Sized {
-    //! Potential assembler
-    //!
-    //! Assemble potential operators by processing batches of cells in parallel
-
-    /// Scalar type
-    type T: RlstScalar + MatrixInverse;
-    /// Number of derivatives
-    const DERIV_SIZE: usize;
-
-    /// Get assembler options
-    fn options(&self) -> &PotentialAssemblerOptions;
-
-    /// Get mutable assembler options
-    fn options_mut(&mut self) -> &mut PotentialAssemblerOptions;
-
-    /// Set (non-singular) quadrature degree for a cell type
-    fn quadrature_degree(&mut self, cell: ReferenceCellType, degree: usize) {
-        *self
-            .options_mut()
-            .quadrature_degrees
-            .get_mut(&cell)
-            .unwrap() = degree;
-    }
-
-    /// Set the maximum size of a batch of cells to send to an assembly function
-    fn batch_size(&mut self, size: usize) {
-        self.options_mut().batch_size = size;
-    }
-
-    /// Return the kernel value to use in the integrand
-    ///
-    /// # Safety
-    /// This method is unsafe to allow `get_unchecked` to be used
-    unsafe fn kernel_value(
-        &self,
-        k: &RlstArray<Self::T, 3>,
-        normals: &RlstArray<<Self::T as RlstScalar>::Real, 2>,
-        index: usize,
-        point_index: usize,
-    ) -> Self::T;
-
-    /// Evaluate the kernel values for all sources and all targets
-    ///
-    /// For every source, the kernel is evaluated for every target.
-    fn kernel_assemble_st(
-        &self,
-        sources: &[<Self::T as RlstScalar>::Real],
-        targets: &[<Self::T as RlstScalar>::Real],
-        result: &mut [Self::T],
-    );
-
-    /// Assemble into a dense matrix
-    fn assemble_into_dense<
-        G: Grid<T = <Self::T as RlstScalar>::Real, EntityDescriptor = ReferenceCellType> + Sync,
-        Element: FiniteElement<T = Self::T> + Sync,
-    >(
-        &self,
-        output: &mut RlstArray<Self::T, 2>,
-        space: &(impl FunctionSpace<Grid = G, FiniteElement = Element> + Sync),
-        points: &RlstArray<<Self::T as RlstScalar>::Real, 2>,
-    ) {
-        if !space.is_serial() {
-            panic!("Dense assembly can only be used for function spaces stored in serial");
-        }
-        if output.shape()[0] != points.shape()[1] || output.shape()[1] != space.global_size() {
-            panic!("Matrix has wrong shape");
-        }
-
-        let colouring = space.cell_colouring();
-
-        let batch_size = self.options().batch_size;
-
-        for cell_type in space.grid().entity_types(2) {
-            let npts = self.options().quadrature_degrees[cell_type];
-            let qrule = simplex_rule(*cell_type, npts).unwrap();
-            let mut qpoints = rlst_dynamic_array2!(<Self::T as RlstScalar>::Real, [2, npts]);
-            for i in 0..npts {
-                for j in 0..2 {
-                    *qpoints.get_mut([j, i]).unwrap() =
-                        num::cast::<f64, <Self::T as RlstScalar>::Real>(qrule.points[2 * i + j])
-                            .unwrap();
-                }
-            }
-            let qweights = qrule
-                .weights
-                .iter()
-                .map(|w| num::cast::<f64, <Self::T as RlstScalar>::Real>(*w).unwrap())
-                .collect::<Vec<_>>();
-
-            let element = space.element(*cell_type);
-            let mut table = rlst_dynamic_array4!(Self::T, element.tabulate_array_shape(0, npts));
-            element.tabulate(&qpoints, 0, &mut table);
-
-            let output_raw = RawData2D {
-                data: output.data_mut().as_mut_ptr(),
-                shape: output.shape(),
-            };
-
-            for c in &colouring[cell_type] {
-                let mut cells: Vec<&[usize]> = vec![];
-
-                let mut start = 0;
-                while start < c.len() {
-                    let end = if start + batch_size < c.len() {
-                        start + batch_size
-                    } else {
-                        c.len()
-                    };
-
-                    cells.push(&c[start..end]);
-                    start = end
-                }
-
-                let numtasks = cells.len();
-                let r: usize = (0..numtasks)
-                    .into_par_iter()
-                    .map(&|t| {
-                        assemble_batch::<Self::T, G, Element>(
-                            self,
-                            Self::DERIV_SIZE,
-                            &output_raw,
-                            *cell_type,
-                            space,
-                            points,
-                            cells[t],
-                            &qpoints,
-                            &qweights,
-                            &table,
-                        )
-                    })
-                    .sum();
-                assert_eq!(r, numtasks);
-            }
-        }
-    }
-}
+pub use assemblers::PotentialAssembler;
diff --git a/src/assembly/potential/assemblers.rs b/src/assembly/potential/assemblers.rs
new file mode 100644
index 00000000..db0f4281
--- /dev/null
+++ b/src/assembly/potential/assemblers.rs
@@ -0,0 +1,252 @@
+//! Potential assemblers
+mod double_layer;
+mod single_layer;
+
+use crate::assembly::common::{RawData2D, RlstArray};
+use crate::assembly::potential::cell_assemblers::PotentialCellAssembler;
+use crate::quadrature::simplex_rules::simplex_rule;
+use crate::traits::{
+    CellAssembler, FunctionSpace, KernelEvaluator, PotentialAssembly, PotentialIntegrand,
+};
+use itertools::izip;
+use ndelement::traits::FiniteElement;
+use ndelement::types::ReferenceCellType;
+use ndgrid::traits::Grid;
+use rayon::prelude::*;
+use rlst::{
+    rlst_dynamic_array2, rlst_dynamic_array4, DefaultIterator, MatrixInverse, RandomAccessMut,
+    RawAccess, RawAccessMut, RlstScalar, Shape,
+};
+use std::collections::HashMap;
+
+/// Assemble the contribution to the terms of a matrix for a batch of non-adjacent cells
+#[allow(clippy::too_many_arguments)]
+fn assemble_batch<
+    T: RlstScalar + MatrixInverse,
+    Space: FunctionSpace<T = T> + Sync,
+    Integrand: PotentialIntegrand<T = T>,
+    Kernel: KernelEvaluator<T = T>,
+>(
+    assembler: &PotentialAssembler<T, Integrand, Kernel>,
+    deriv_size: usize,
+    output: &RawData2D<T>,
+    cell_type: ReferenceCellType,
+    space: &Space,
+    evaluation_points: &RlstArray<T::Real, 2>,
+    cells: &[usize],
+    points: &RlstArray<T::Real, 2>,
+    weights: &[T::Real],
+    table: &RlstArray<T, 4>,
+) -> usize {
+    let npts = weights.len();
+    let nevalpts = evaluation_points.shape()[1];
+    debug_assert!(points.shape()[1] == npts);
+
+    let grid = space.grid();
+
+    assert_eq!(grid.geometry_dim(), 3);
+    assert_eq!(grid.topology_dim(), 2);
+
+    let evaluator = grid.geometry_map(cell_type, points.data());
+    let mut a = PotentialCellAssembler::new(
+        npts,
+        nevalpts,
+        deriv_size,
+        &assembler.integrand,
+        &assembler.kernel,
+        evaluator,
+        table,
+        evaluation_points,
+        weights,
+    );
+
+    let mut local_mat = rlst_dynamic_array2!(T, [nevalpts, space.element(cell_type).dim()]);
+
+    for cell in cells {
+        a.set_cell(*cell);
+        a.assemble(&mut local_mat);
+
+        let dofs = space.cell_dofs(*cell).unwrap();
+        for (dof, col) in izip!(dofs, local_mat.col_iter()) {
+            for (eval_index, entry) in col.iter().enumerate() {
+                unsafe {
+                    *output.data.add(eval_index + output.shape[0] * *dof) += entry;
+                }
+            }
+        }
+    }
+    1
+}
+
+/// Options for a potential assembler
+pub struct PotentialAssemblerOptions {
+    /// Number of points used in quadrature for non-singular integrals
+    quadrature_degrees: HashMap<ReferenceCellType, usize>,
+    /// Maximum size of each batch of cells to send to an assembly function
+    batch_size: usize,
+}
+
+impl Default for PotentialAssemblerOptions {
+    fn default() -> Self {
+        use ReferenceCellType::{Quadrilateral, Triangle};
+        Self {
+            quadrature_degrees: HashMap::from([(Triangle, 37), (Quadrilateral, 37)]),
+            batch_size: 128,
+        }
+    }
+}
+
+/// Potential assembler
+///
+/// Assemble potential operators by processing batches of cells in parallel
+pub struct PotentialAssembler<
+    T: RlstScalar + MatrixInverse,
+    Integrand: PotentialIntegrand<T = T>,
+    Kernel: KernelEvaluator<T = T>,
+> {
+    pub(crate) integrand: Integrand,
+    pub(crate) kernel: Kernel,
+    pub(crate) options: PotentialAssemblerOptions,
+    pub(crate) deriv_size: usize,
+}
+
+unsafe impl<
+        T: RlstScalar + MatrixInverse,
+        Integrand: PotentialIntegrand<T = T>,
+        Kernel: KernelEvaluator<T = T>,
+    > Sync for PotentialAssembler<T, Integrand, Kernel>
+{
+}
+
+impl<
+        T: RlstScalar + MatrixInverse,
+        Integrand: PotentialIntegrand<T = T>,
+        Kernel: KernelEvaluator<T = T>,
+    > PotentialAssembler<T, Integrand, Kernel>
+{
+    /// Create new
+    fn new(integrand: Integrand, kernel: Kernel, deriv_size: usize) -> Self {
+        Self {
+            integrand,
+            kernel,
+            options: PotentialAssemblerOptions::default(),
+            deriv_size,
+        }
+    }
+
+    /// Set (non-singular) quadrature degree for a cell type
+    pub fn quadrature_degree(&mut self, cell: ReferenceCellType, degree: usize) {
+        *self.options.quadrature_degrees.get_mut(&cell).unwrap() = degree;
+    }
+
+    /// Set the maximum size of a batch of cells to send to an assembly function
+    pub fn batch_size(&mut self, size: usize) {
+        self.options.batch_size = size;
+    }
+
+    fn assemble<Space: FunctionSpace<T = T> + Sync>(
+        &self,
+        output: &RawData2D<T>,
+        space: &Space,
+        points: &RlstArray<T::Real, 2>,
+        colouring: &HashMap<ReferenceCellType, Vec<Vec<usize>>>,
+    ) {
+        if !space.is_serial() {
+            panic!("Dense assembly can only be used for function spaces stored in serial");
+        }
+        if output.shape[0] != points.shape()[1] || output.shape[1] != space.global_size() {
+            panic!("Matrix has wrong shape");
+        }
+
+        let batch_size = self.options.batch_size;
+
+        for cell_type in space.grid().entity_types(2) {
+            let npts = self.options.quadrature_degrees[cell_type];
+            let qrule = simplex_rule(*cell_type, npts).unwrap();
+            let mut qpoints = rlst_dynamic_array2!(T::Real, [2, npts]);
+            for i in 0..npts {
+                for j in 0..2 {
+                    *qpoints.get_mut([j, i]).unwrap() =
+                        num::cast::<f64, T::Real>(qrule.points[2 * i + j]).unwrap();
+                }
+            }
+            let qweights = qrule
+                .weights
+                .iter()
+                .map(|w| num::cast::<f64, T::Real>(*w).unwrap())
+                .collect::<Vec<_>>();
+
+            let element = space.element(*cell_type);
+            let mut table = rlst_dynamic_array4!(T, element.tabulate_array_shape(0, npts));
+            element.tabulate(&qpoints, 0, &mut table);
+
+            for c in &colouring[cell_type] {
+                let mut cells: Vec<&[usize]> = vec![];
+
+                let mut start = 0;
+                while start < c.len() {
+                    let end = if start + batch_size < c.len() {
+                        start + batch_size
+                    } else {
+                        c.len()
+                    };
+
+                    cells.push(&c[start..end]);
+                    start = end
+                }
+
+                let numtasks = cells.len();
+                let r: usize = (0..numtasks)
+                    .into_par_iter()
+                    .map(&|t| {
+                        assemble_batch(
+                            self,
+                            self.deriv_size,
+                            output,
+                            *cell_type,
+                            space,
+                            points,
+                            cells[t],
+                            &qpoints,
+                            &qweights,
+                            &table,
+                        )
+                    })
+                    .sum();
+                assert_eq!(r, numtasks);
+            }
+        }
+    }
+}
+
+impl<
+        T: RlstScalar + MatrixInverse,
+        Integrand: PotentialIntegrand<T = T>,
+        Kernel: KernelEvaluator<T = T>,
+    > PotentialAssembly for PotentialAssembler<T, Integrand, Kernel>
+{
+    type T = T;
+
+    fn assemble_into_dense<Space: FunctionSpace<T = T> + Sync>(
+        &self,
+        output: &mut RlstArray<T, 2>,
+        space: &Space,
+        points: &RlstArray<T::Real, 2>,
+    ) {
+        if !space.is_serial() {
+            panic!("Dense assembly can only be used for function spaces stored in serial");
+        }
+        if output.shape()[0] != points.shape()[1] || output.shape()[1] != space.global_size() {
+            panic!("Matrix has wrong shape");
+        }
+
+        let output_raw = RawData2D {
+            data: output.data_mut().as_mut_ptr(),
+            shape: output.shape(),
+        };
+
+        let colouring = space.cell_colouring();
+
+        self.assemble(&output_raw, space, points, &colouring);
+    }
+}
diff --git a/src/assembly/potential/assemblers/double_layer.rs b/src/assembly/potential/assemblers/double_layer.rs
new file mode 100644
index 00000000..673ce969
--- /dev/null
+++ b/src/assembly/potential/assemblers/double_layer.rs
@@ -0,0 +1,42 @@
+//! Double layer assemblers
+use super::PotentialAssembler;
+use crate::assembly::{
+    common::GreenKernelEvalType, kernels::KernelEvaluator,
+    potential::integrands::DoubleLayerPotentialIntegrand,
+};
+use green_kernels::{helmholtz_3d::Helmholtz3dKernel, laplace_3d::Laplace3dKernel, traits::Kernel};
+use rlst::{MatrixInverse, RlstScalar};
+
+impl<T: RlstScalar + MatrixInverse, K: Kernel<T = T>>
+    PotentialAssembler<T, DoubleLayerPotentialIntegrand<T>, KernelEvaluator<T, K>>
+{
+    /// Create a new Double layer assembler
+    pub fn new_double_layer(kernel: KernelEvaluator<T, K>) -> Self {
+        Self::new(DoubleLayerPotentialIntegrand::new(), kernel, 4)
+    }
+}
+impl<T: RlstScalar + MatrixInverse>
+    PotentialAssembler<T, DoubleLayerPotentialIntegrand<T>, KernelEvaluator<T, Laplace3dKernel<T>>>
+{
+    /// Create a new Laplace Double layer assembler
+    pub fn new_laplace_double_layer() -> Self {
+        Self::new_double_layer(KernelEvaluator::new_laplace(
+            GreenKernelEvalType::ValueDeriv,
+        ))
+    }
+}
+impl<T: RlstScalar<Complex = T> + MatrixInverse>
+    PotentialAssembler<
+        T,
+        DoubleLayerPotentialIntegrand<T>,
+        KernelEvaluator<T, Helmholtz3dKernel<T>>,
+    >
+{
+    /// Create a new Helmholtz Double layer assembler
+    pub fn new_helmholtz_double_layer(wavenumber: T::Real) -> Self {
+        Self::new_double_layer(KernelEvaluator::new_helmholtz(
+            wavenumber,
+            GreenKernelEvalType::ValueDeriv,
+        ))
+    }
+}
diff --git a/src/assembly/potential/assemblers/single_layer.rs b/src/assembly/potential/assemblers/single_layer.rs
new file mode 100644
index 00000000..aa5cb39e
--- /dev/null
+++ b/src/assembly/potential/assemblers/single_layer.rs
@@ -0,0 +1,40 @@
+//! Single layer assemblers
+use super::PotentialAssembler;
+use crate::assembly::{
+    common::GreenKernelEvalType, kernels::KernelEvaluator,
+    potential::integrands::SingleLayerPotentialIntegrand,
+};
+use green_kernels::{helmholtz_3d::Helmholtz3dKernel, laplace_3d::Laplace3dKernel, traits::Kernel};
+use rlst::{MatrixInverse, RlstScalar};
+
+impl<T: RlstScalar + MatrixInverse, K: Kernel<T = T>>
+    PotentialAssembler<T, SingleLayerPotentialIntegrand<T>, KernelEvaluator<T, K>>
+{
+    /// Create a new single layer assembler
+    pub fn new_single_layer(kernel: KernelEvaluator<T, K>) -> Self {
+        Self::new(SingleLayerPotentialIntegrand::new(), kernel, 1)
+    }
+}
+impl<T: RlstScalar + MatrixInverse>
+    PotentialAssembler<T, SingleLayerPotentialIntegrand<T>, KernelEvaluator<T, Laplace3dKernel<T>>>
+{
+    /// Create a new Laplace single layer assembler
+    pub fn new_laplace_single_layer() -> Self {
+        Self::new_single_layer(KernelEvaluator::new_laplace(GreenKernelEvalType::Value))
+    }
+}
+impl<T: RlstScalar<Complex = T> + MatrixInverse>
+    PotentialAssembler<
+        T,
+        SingleLayerPotentialIntegrand<T>,
+        KernelEvaluator<T, Helmholtz3dKernel<T>>,
+    >
+{
+    /// Create a new Helmholtz single layer assembler
+    pub fn new_helmholtz_single_layer(wavenumber: T::Real) -> Self {
+        Self::new_single_layer(KernelEvaluator::new_helmholtz(
+            wavenumber,
+            GreenKernelEvalType::Value,
+        ))
+    }
+}
diff --git a/src/assembly/potential/cell_assemblers.rs b/src/assembly/potential/cell_assemblers.rs
new file mode 100644
index 00000000..199f1376
--- /dev/null
+++ b/src/assembly/potential/cell_assemblers.rs
@@ -0,0 +1,112 @@
+//! Assemblers that assemble the contributions to the global matrix due to a single cell
+
+use crate::assembly::common::{AssemblerGeometry, RlstArray};
+use crate::traits::{CellAssembler, KernelEvaluator, PotentialIntegrand};
+use ndgrid::traits::GeometryMap;
+use num::Zero;
+use rlst::{
+    rlst_dynamic_array2, rlst_dynamic_array3, DefaultIteratorMut, RawAccess, RawAccessMut,
+    RlstScalar,
+};
+
+pub struct PotentialCellAssembler<
+    'a,
+    T: RlstScalar,
+    I: PotentialIntegrand<T = T>,
+    G: GeometryMap<T = T::Real>,
+    K: KernelEvaluator<T = T>,
+> {
+    integrand: &'a I,
+    kernel: &'a K,
+    evaluator: G,
+    table: &'a RlstArray<T, 4>,
+    k: RlstArray<T, 3>,
+    mapped_pts: RlstArray<T::Real, 2>,
+    normals: RlstArray<T::Real, 2>,
+    jacobians: RlstArray<T::Real, 2>,
+    jdet: Vec<T::Real>,
+    evaluation_points: &'a RlstArray<T::Real, 2>,
+    weights: &'a [T::Real],
+    cell: usize,
+}
+
+impl<
+        'a,
+        T: RlstScalar,
+        I: PotentialIntegrand<T = T>,
+        G: GeometryMap<T = T::Real>,
+        K: KernelEvaluator<T = T>,
+    > PotentialCellAssembler<'a, T, I, G, K>
+{
+    #[allow(clippy::too_many_arguments)]
+    pub fn new(
+        npts: usize,
+        nevalpts: usize,
+        deriv_size: usize,
+        integrand: &'a I,
+        kernel: &'a K,
+        evaluator: G,
+        table: &'a RlstArray<T, 4>,
+        evaluation_points: &'a RlstArray<T::Real, 2>,
+        weights: &'a [T::Real],
+    ) -> Self {
+        Self {
+            integrand,
+            kernel,
+            evaluator,
+            table,
+            k: rlst_dynamic_array3!(T, [deriv_size, npts, nevalpts]),
+            mapped_pts: rlst_dynamic_array2!(T::Real, [3, npts]),
+            normals: rlst_dynamic_array2!(T::Real, [3, npts]),
+            jacobians: rlst_dynamic_array2!(T::Real, [6, npts]),
+            jdet: vec![T::Real::zero(); npts],
+            evaluation_points,
+            weights,
+            cell: 0,
+        }
+    }
+}
+impl<
+        'a,
+        T: RlstScalar,
+        I: PotentialIntegrand<T = T>,
+        G: GeometryMap<T = T::Real>,
+        K: KernelEvaluator<T = T>,
+    > CellAssembler for PotentialCellAssembler<'a, T, I, G, K>
+{
+    type T = T;
+    fn set_cell(&mut self, cell: usize) {
+        self.cell = cell;
+        self.evaluator.points(cell, self.mapped_pts.data_mut());
+        self.evaluator.jacobians_dets_normals(
+            cell,
+            self.jacobians.data_mut(),
+            &mut self.jdet,
+            self.normals.data_mut(),
+        );
+    }
+
+    fn assemble(&mut self, local_mat: &mut RlstArray<T, 2>) {
+        self.kernel.assemble_st(
+            self.mapped_pts.data(),
+            self.evaluation_points.data(),
+            self.k.data_mut(),
+        );
+
+        let geometry =
+            AssemblerGeometry::new(&self.mapped_pts, &self.normals, &self.jacobians, &self.jdet);
+
+        for (i, mut col) in local_mat.col_iter_mut().enumerate() {
+            for (eval_index, entry) in col.iter_mut().enumerate() {
+                *entry = T::zero();
+                for (index, wt) in self.weights.iter().enumerate() {
+                    *entry += unsafe {
+                        self.integrand
+                            .evaluate(self.table, index, eval_index, i, &self.k, &geometry)
+                            * num::cast::<T::Real, T>(*wt * self.jdet[index]).unwrap()
+                    };
+                }
+            }
+        }
+    }
+}
diff --git a/src/assembly/potential/double_layer.rs b/src/assembly/potential/double_layer.rs
deleted file mode 100644
index 0ae223da..00000000
--- a/src/assembly/potential/double_layer.rs
+++ /dev/null
@@ -1,73 +0,0 @@
-//! Double layer potential assemblers
-use super::{PotentialAssembler, PotentialAssemblerOptions};
-use crate::assembly::common::{GreenKernelEvalType, RlstArray};
-use green_kernels::{helmholtz_3d::Helmholtz3dKernel, laplace_3d::Laplace3dKernel, traits::Kernel};
-use rlst::{MatrixInverse, RlstScalar, UnsafeRandomAccessByRef};
-
-/// Assembler for a double layer potential operator
-pub struct DoubleLayerPotentialAssembler<T: RlstScalar + MatrixInverse, K: Kernel<T = T>> {
-    kernel: K,
-    options: PotentialAssemblerOptions,
-}
-impl<T: RlstScalar + MatrixInverse, K: Kernel<T = T>> DoubleLayerPotentialAssembler<T, K> {
-    /// Create a new double layer potential assembler
-    pub fn new(kernel: K) -> Self {
-        Self {
-            kernel,
-            options: PotentialAssemblerOptions::default(),
-        }
-    }
-}
-impl<T: RlstScalar + MatrixInverse> DoubleLayerPotentialAssembler<T, Laplace3dKernel<T>> {
-    /// Create a new Laplace double layer potential assembler
-    pub fn new_laplace() -> Self {
-        Self::new(Laplace3dKernel::<T>::new())
-    }
-}
-impl<T: RlstScalar<Complex = T> + MatrixInverse>
-    DoubleLayerPotentialAssembler<T, Helmholtz3dKernel<T>>
-{
-    /// Create a new Helmholtz double layer potential assembler
-    pub fn new_helmholtz(wavenumber: T::Real) -> Self {
-        Self::new(Helmholtz3dKernel::<T>::new(wavenumber))
-    }
-}
-
-impl<T: RlstScalar + MatrixInverse, K: Kernel<T = T>> PotentialAssembler
-    for DoubleLayerPotentialAssembler<T, K>
-{
-    const DERIV_SIZE: usize = 4;
-    type T = T;
-
-    fn options(&self) -> &PotentialAssemblerOptions {
-        &self.options
-    }
-    fn options_mut(&mut self) -> &mut PotentialAssemblerOptions {
-        &mut self.options
-    }
-
-    unsafe fn kernel_value(
-        &self,
-        k: &RlstArray<T, 3>,
-        normals: &RlstArray<T::Real, 2>,
-        index: usize,
-        point_index: usize,
-    ) -> T {
-        -*k.get_unchecked([1, index, point_index])
-            * num::cast::<T::Real, T>(*normals.get_unchecked([0, index])).unwrap()
-            - *k.get_unchecked([2, index, point_index])
-                * num::cast::<T::Real, T>(*normals.get_unchecked([1, index])).unwrap()
-            - *k.get_unchecked([3, index, point_index])
-                * num::cast::<T::Real, T>(*normals.get_unchecked([2, index])).unwrap()
-    }
-
-    fn kernel_assemble_st(
-        &self,
-        sources: &[<Self::T as RlstScalar>::Real],
-        targets: &[<Self::T as RlstScalar>::Real],
-        result: &mut [Self::T],
-    ) {
-        self.kernel
-            .assemble_st(GreenKernelEvalType::ValueDeriv, sources, targets, result);
-    }
-}
diff --git a/src/assembly/potential/integrands.rs b/src/assembly/potential/integrands.rs
new file mode 100644
index 00000000..5f6fd3e1
--- /dev/null
+++ b/src/assembly/potential/integrands.rs
@@ -0,0 +1,6 @@
+//! Integrands
+mod double_layer;
+mod single_layer;
+
+pub use double_layer::DoubleLayerPotentialIntegrand;
+pub use single_layer::SingleLayerPotentialIntegrand;
diff --git a/src/assembly/potential/integrands/double_layer.rs b/src/assembly/potential/integrands/double_layer.rs
new file mode 100644
index 00000000..2f0b6cd9
--- /dev/null
+++ b/src/assembly/potential/integrands/double_layer.rs
@@ -0,0 +1,46 @@
+//! Double layer integrand
+use crate::assembly::common::RlstArray;
+use crate::traits::{CellGeometry, PotentialIntegrand};
+use rlst::{RlstScalar, UnsafeRandomAccessByRef};
+
+pub struct DoubleLayerPotentialIntegrand<T: RlstScalar> {
+    _t: std::marker::PhantomData<T>,
+}
+
+impl<T: RlstScalar> DoubleLayerPotentialIntegrand<T> {
+    pub fn new() -> Self {
+        Self {
+            _t: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<T: RlstScalar> PotentialIntegrand for DoubleLayerPotentialIntegrand<T> {
+    type T = T;
+
+    unsafe fn evaluate(
+        &self,
+        table: &RlstArray<T, 4>,
+        point_index: usize,
+        eval_index: usize,
+        basis_index: usize,
+        k: &RlstArray<T, 3>,
+        geometry: &impl CellGeometry<T = T::Real>,
+    ) -> T {
+        -(*k.get_unchecked([1, point_index, eval_index])
+            * num::cast::<T::Real, T>(*geometry.normals().get_unchecked([0, point_index])).unwrap()
+            + *k.get_unchecked([2, point_index, eval_index])
+                * num::cast::<T::Real, T>(*geometry.normals().get_unchecked([1, point_index]))
+                    .unwrap()
+            + *k.get_unchecked([3, point_index, eval_index])
+                * num::cast::<T::Real, T>(*geometry.normals().get_unchecked([2, point_index]))
+                    .unwrap())
+            * *table.get_unchecked([0, point_index, basis_index, 0])
+    }
+}
+
+impl<T: RlstScalar> Default for DoubleLayerPotentialIntegrand<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
diff --git a/src/assembly/potential/integrands/single_layer.rs b/src/assembly/potential/integrands/single_layer.rs
new file mode 100644
index 00000000..7a1a28a2
--- /dev/null
+++ b/src/assembly/potential/integrands/single_layer.rs
@@ -0,0 +1,39 @@
+//! Single layer integrand
+use crate::assembly::common::RlstArray;
+use crate::traits::{CellGeometry, PotentialIntegrand};
+use rlst::{RlstScalar, UnsafeRandomAccessByRef};
+
+pub struct SingleLayerPotentialIntegrand<T: RlstScalar> {
+    _t: std::marker::PhantomData<T>,
+}
+
+impl<T: RlstScalar> SingleLayerPotentialIntegrand<T> {
+    pub fn new() -> Self {
+        Self {
+            _t: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<T: RlstScalar> PotentialIntegrand for SingleLayerPotentialIntegrand<T> {
+    type T = T;
+
+    unsafe fn evaluate(
+        &self,
+        table: &RlstArray<T, 4>,
+        point_index: usize,
+        eval_index: usize,
+        basis_index: usize,
+        k: &RlstArray<T, 3>,
+        _geometry: &impl CellGeometry<T = T::Real>,
+    ) -> T {
+        *k.get_unchecked([0, point_index, eval_index])
+            * *table.get_unchecked([0, point_index, basis_index, 0])
+    }
+}
+
+impl<T: RlstScalar> Default for SingleLayerPotentialIntegrand<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
diff --git a/src/assembly/potential/single_layer.rs b/src/assembly/potential/single_layer.rs
deleted file mode 100644
index 657f04b3..00000000
--- a/src/assembly/potential/single_layer.rs
+++ /dev/null
@@ -1,68 +0,0 @@
-//! Single layer potential assemblers
-use super::{PotentialAssembler, PotentialAssemblerOptions};
-use crate::assembly::common::{GreenKernelEvalType, RlstArray};
-use green_kernels::{helmholtz_3d::Helmholtz3dKernel, laplace_3d::Laplace3dKernel, traits::Kernel};
-use rlst::{MatrixInverse, RlstScalar, UnsafeRandomAccessByRef};
-
-/// Assembler for a single layer potential operator
-pub struct SingleLayerPotentialAssembler<T: RlstScalar + MatrixInverse, K: Kernel<T = T>> {
-    kernel: K,
-    options: PotentialAssemblerOptions,
-}
-impl<T: RlstScalar + MatrixInverse, K: Kernel<T = T>> SingleLayerPotentialAssembler<T, K> {
-    /// Create a new single layer potential assembler
-    pub fn new(kernel: K) -> Self {
-        Self {
-            kernel,
-            options: PotentialAssemblerOptions::default(),
-        }
-    }
-}
-impl<T: RlstScalar + MatrixInverse> SingleLayerPotentialAssembler<T, Laplace3dKernel<T>> {
-    /// Create a new Laplace single layer potential assembler
-    pub fn new_laplace() -> Self {
-        Self::new(Laplace3dKernel::<T>::new())
-    }
-}
-impl<T: RlstScalar<Complex = T> + MatrixInverse>
-    SingleLayerPotentialAssembler<T, Helmholtz3dKernel<T>>
-{
-    /// Create a new Helmholtz single layer potential assembler
-    pub fn new_helmholtz(wavenumber: T::Real) -> Self {
-        Self::new(Helmholtz3dKernel::<T>::new(wavenumber))
-    }
-}
-
-impl<T: RlstScalar + MatrixInverse, K: Kernel<T = T>> PotentialAssembler
-    for SingleLayerPotentialAssembler<T, K>
-{
-    const DERIV_SIZE: usize = 1;
-    type T = T;
-
-    fn options(&self) -> &PotentialAssemblerOptions {
-        &self.options
-    }
-    fn options_mut(&mut self) -> &mut PotentialAssemblerOptions {
-        &mut self.options
-    }
-
-    unsafe fn kernel_value(
-        &self,
-        k: &RlstArray<T, 3>,
-        _normals: &RlstArray<T::Real, 2>,
-        index: usize,
-        point_index: usize,
-    ) -> T {
-        *k.get_unchecked([0, index, point_index])
-    }
-
-    fn kernel_assemble_st(
-        &self,
-        sources: &[<Self::T as RlstScalar>::Real],
-        targets: &[<Self::T as RlstScalar>::Real],
-        result: &mut [Self::T],
-    ) {
-        self.kernel
-            .assemble_st(GreenKernelEvalType::Value, sources, targets, result);
-    }
-}
diff --git a/src/traits.rs b/src/traits.rs
index 5962fb4b..57224fce 100644
--- a/src/traits.rs
+++ b/src/traits.rs
@@ -6,7 +6,8 @@ mod function;
 #[cfg(feature = "mpi")]
 pub use assembly::ParallelBoundaryAssembly;
 pub use assembly::{
-    BoundaryAssembly, BoundaryIntegrand, CellGeometry, CellPairAssembler, KernelEvaluator,
+    BoundaryAssembly, BoundaryIntegrand, CellAssembler, CellGeometry, CellPairAssembler,
+    KernelEvaluator, PotentialAssembly, PotentialIntegrand,
 };
 pub use function::FunctionSpace;
 #[cfg(feature = "mpi")]
diff --git a/src/traits/assembly.rs b/src/traits/assembly.rs
index ca8ff204..d48542dd 100644
--- a/src/traits/assembly.rs
+++ b/src/traits/assembly.rs
@@ -1,14 +1,13 @@
 //! Assembly
+mod boundary;
+mod potential;
 use crate::assembly::common::RlstArray;
-use crate::traits::FunctionSpace;
-use ndelement::types::ReferenceCellType;
-use rlst::{CsrMatrix, RlstScalar};
-use std::collections::HashMap;
+use rlst::RlstScalar;
 
 #[cfg(feature = "mpi")]
-use crate::traits::ParallelFunctionSpace;
-#[cfg(feature = "mpi")]
-use mpi::traits::Communicator;
+pub use boundary::ParallelBoundaryAssembly;
+pub use boundary::{BoundaryAssembly, BoundaryIntegrand, CellPairAssembler};
+pub use potential::{CellAssembler, PotentialAssembly, PotentialIntegrand};
 
 pub trait CellGeometry {
     //! Cell geometry
@@ -24,47 +23,6 @@ pub trait CellGeometry {
     fn jdets(&self) -> &[Self::T];
 }
 
-pub trait BoundaryIntegrand {
-    //! Integrand
-    /// Scalar type
-    type T: RlstScalar;
-
-    #[allow(clippy::too_many_arguments)]
-    /// Evaluate integrand for a singular quadrature rule
-    ///
-    /// # Safety
-    /// This method is unsafe to allow `get_unchecked` to be used
-    unsafe fn evaluate_nonsingular(
-        &self,
-        test_table: &RlstArray<Self::T, 4>,
-        trial_table: &RlstArray<Self::T, 4>,
-        test_point_index: usize,
-        trial_point_index: usize,
-        test_basis_index: usize,
-        trial_basis_index: usize,
-        k: &RlstArray<Self::T, 3>,
-        test_geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
-        trial_geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
-    ) -> Self::T;
-
-    #[allow(clippy::too_many_arguments)]
-    /// Evaluate integrand for a non-singular quadrature rule
-    ///
-    /// # Safety
-    /// This method is unsafe to allow `get_unchecked` to be used
-    unsafe fn evaluate_singular(
-        &self,
-        test_table: &RlstArray<Self::T, 4>,
-        trial_table: &RlstArray<Self::T, 4>,
-        point_index: usize,
-        test_basis_index: usize,
-        trial_basis_index: usize,
-        k: &RlstArray<Self::T, 2>,
-        test_geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
-        trial_geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
-    ) -> Self::T;
-}
-
 pub trait KernelEvaluator {
     //! Kernel evaluator
     /// Scalar type
@@ -90,99 +48,3 @@ pub trait KernelEvaluator {
         result: &mut [Self::T],
     );
 }
-
-pub trait CellPairAssembler {
-    //! Assembler for the contributions from a pair of cells
-    /// Scalar type
-    type T: RlstScalar;
-
-    /// Assemble contributions into `local_mat`
-    fn assemble(&mut self, local_mat: &mut RlstArray<Self::T, 2>);
-    /// Set the test cell
-    fn set_test_cell(&mut self, test_cell: usize);
-    /// Set the trial cell
-    fn set_trial_cell(&mut self, trial_cell: usize);
-}
-
-pub trait BoundaryAssembly {
-    //! Functions for boundary assembly
-    /// Scalar type
-    type T: RlstScalar;
-
-    /// Assemble the singular contributions into a dense matrix
-    fn assemble_singular_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
-        &self,
-        output: &mut RlstArray<Self::T, 2>,
-        trial_space: &Space,
-        test_space: &Space,
-    );
-
-    /// Assemble the singular contributions into a CSR sparse matrix
-    fn assemble_singular_into_csr<Space: FunctionSpace<T = Self::T> + Sync>(
-        &self,
-        trial_space: &Space,
-        test_space: &Space,
-    ) -> CsrMatrix<Self::T>;
-
-    /// Assemble the singular correction into a dense matrix
-    ///
-    /// The singular correction is the contribution is the terms for adjacent cells are assembled using an (incorrect) non-singular quadrature rule
-    fn assemble_singular_correction_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
-        &self,
-        output: &mut RlstArray<Self::T, 2>,
-        trial_space: &Space,
-        test_space: &Space,
-    );
-
-    /// Assemble the singular correction into a CSR matrix
-    ///
-    /// The singular correction is the contribution is the terms for adjacent cells are assembled using an (incorrect) non-singular quadrature rule
-    fn assemble_singular_correction_into_csr<Space: FunctionSpace<T = Self::T> + Sync>(
-        &self,
-        trial_space: &Space,
-        test_space: &Space,
-    ) -> CsrMatrix<Self::T>;
-
-    /// Assemble into a dense matrix
-    fn assemble_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
-        &self,
-        output: &mut RlstArray<Self::T, 2>,
-        trial_space: &Space,
-        test_space: &Space,
-    );
-
-    /// Assemble the non-singular contributions into a dense matrix
-    fn assemble_nonsingular_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
-        &self,
-        output: &mut RlstArray<Self::T, 2>,
-        trial_space: &Space,
-        test_space: &Space,
-        trial_colouring: &HashMap<ReferenceCellType, Vec<Vec<usize>>>,
-        test_colouring: &HashMap<ReferenceCellType, Vec<Vec<usize>>>,
-    );
-}
-
-#[cfg(feature = "mpi")]
-pub trait ParallelBoundaryAssembly: BoundaryAssembly {
-    //! Functions for parallel boundary assembly
-
-    /// Assemble the singular contributions into a CSR sparse matrix, indexed by global DOF numbers
-    fn parallel_assemble_singular_into_csr<
-        C: Communicator,
-        Space: ParallelFunctionSpace<C, T = Self::T>,
-    >(
-        &self,
-        trial_space: &Space,
-        test_space: &Space,
-    ) -> CsrMatrix<Self::T>;
-
-    /// Assemble the singular contributions into a CSR sparse matrix, indexed by global DOF numbers
-    fn parallel_assemble_singular_correction_into_csr<
-        C: Communicator,
-        Space: ParallelFunctionSpace<C, T = Self::T>,
-    >(
-        &self,
-        trial_space: &Space,
-        test_space: &Space,
-    ) -> CsrMatrix<Self::T>;
-}
diff --git a/src/traits/assembly/boundary.rs b/src/traits/assembly/boundary.rs
new file mode 100644
index 00000000..5ab59c9a
--- /dev/null
+++ b/src/traits/assembly/boundary.rs
@@ -0,0 +1,148 @@
+//! Traits for boundary assembly
+use super::CellGeometry;
+use crate::assembly::common::RlstArray;
+use crate::traits::FunctionSpace;
+#[cfg(feature = "mpi")]
+use crate::traits::ParallelFunctionSpace;
+#[cfg(feature = "mpi")]
+use mpi::traits::Communicator;
+use ndelement::types::ReferenceCellType;
+use rlst::{CsrMatrix, RlstScalar};
+use std::collections::HashMap;
+
+pub trait BoundaryIntegrand {
+    //! Integrand
+    /// Scalar type
+    type T: RlstScalar;
+
+    #[allow(clippy::too_many_arguments)]
+    /// Evaluate integrand for a singular quadrature rule
+    ///
+    /// # Safety
+    /// This method is unsafe to allow `get_unchecked` to be used
+    unsafe fn evaluate_nonsingular(
+        &self,
+        test_table: &RlstArray<Self::T, 4>,
+        trial_table: &RlstArray<Self::T, 4>,
+        test_point_index: usize,
+        trial_point_index: usize,
+        test_basis_index: usize,
+        trial_basis_index: usize,
+        k: &RlstArray<Self::T, 3>,
+        test_geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
+        trial_geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
+    ) -> Self::T;
+
+    #[allow(clippy::too_many_arguments)]
+    /// Evaluate integrand for a non-singular quadrature rule
+    ///
+    /// # Safety
+    /// This method is unsafe to allow `get_unchecked` to be used
+    unsafe fn evaluate_singular(
+        &self,
+        test_table: &RlstArray<Self::T, 4>,
+        trial_table: &RlstArray<Self::T, 4>,
+        point_index: usize,
+        test_basis_index: usize,
+        trial_basis_index: usize,
+        k: &RlstArray<Self::T, 2>,
+        test_geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
+        trial_geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
+    ) -> Self::T;
+}
+
+pub trait CellPairAssembler {
+    //! Assembler for the contributions from a pair of cells
+    /// Scalar type
+    type T: RlstScalar;
+
+    /// Assemble contributions into `local_mat`
+    fn assemble(&mut self, local_mat: &mut RlstArray<Self::T, 2>);
+    /// Set the test cell
+    fn set_test_cell(&mut self, test_cell: usize);
+    /// Set the trial cell
+    fn set_trial_cell(&mut self, trial_cell: usize);
+}
+
+pub trait BoundaryAssembly {
+    //! Functions for boundary assembly
+    /// Scalar type
+    type T: RlstScalar;
+
+    /// Assemble the singular contributions into a dense matrix
+    fn assemble_singular_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
+        &self,
+        output: &mut RlstArray<Self::T, 2>,
+        trial_space: &Space,
+        test_space: &Space,
+    );
+
+    /// Assemble the singular contributions into a CSR sparse matrix
+    fn assemble_singular_into_csr<Space: FunctionSpace<T = Self::T> + Sync>(
+        &self,
+        trial_space: &Space,
+        test_space: &Space,
+    ) -> CsrMatrix<Self::T>;
+
+    /// Assemble the singular correction into a dense matrix
+    ///
+    /// The singular correction is the contribution is the terms for adjacent cells are assembled using an (incorrect) non-singular quadrature rule
+    fn assemble_singular_correction_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
+        &self,
+        output: &mut RlstArray<Self::T, 2>,
+        trial_space: &Space,
+        test_space: &Space,
+    );
+
+    /// Assemble the singular correction into a CSR matrix
+    ///
+    /// The singular correction is the contribution is the terms for adjacent cells are assembled using an (incorrect) non-singular quadrature rule
+    fn assemble_singular_correction_into_csr<Space: FunctionSpace<T = Self::T> + Sync>(
+        &self,
+        trial_space: &Space,
+        test_space: &Space,
+    ) -> CsrMatrix<Self::T>;
+
+    /// Assemble into a dense matrix
+    fn assemble_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
+        &self,
+        output: &mut RlstArray<Self::T, 2>,
+        trial_space: &Space,
+        test_space: &Space,
+    );
+
+    /// Assemble the non-singular contributions into a dense matrix
+    fn assemble_nonsingular_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
+        &self,
+        output: &mut RlstArray<Self::T, 2>,
+        trial_space: &Space,
+        test_space: &Space,
+        trial_colouring: &HashMap<ReferenceCellType, Vec<Vec<usize>>>,
+        test_colouring: &HashMap<ReferenceCellType, Vec<Vec<usize>>>,
+    );
+}
+
+#[cfg(feature = "mpi")]
+pub trait ParallelBoundaryAssembly: BoundaryAssembly {
+    //! Functions for parallel boundary assembly
+
+    /// Assemble the singular contributions into a CSR sparse matrix, indexed by global DOF numbers
+    fn parallel_assemble_singular_into_csr<
+        C: Communicator,
+        Space: ParallelFunctionSpace<C, T = Self::T>,
+    >(
+        &self,
+        trial_space: &Space,
+        test_space: &Space,
+    ) -> CsrMatrix<Self::T>;
+
+    /// Assemble the singular contributions into a CSR sparse matrix, indexed by global DOF numbers
+    fn parallel_assemble_singular_correction_into_csr<
+        C: Communicator,
+        Space: ParallelFunctionSpace<C, T = Self::T>,
+    >(
+        &self,
+        trial_space: &Space,
+        test_space: &Space,
+    ) -> CsrMatrix<Self::T>;
+}
diff --git a/src/traits/assembly/potential.rs b/src/traits/assembly/potential.rs
new file mode 100644
index 00000000..775efabd
--- /dev/null
+++ b/src/traits/assembly/potential.rs
@@ -0,0 +1,50 @@
+//! Traits for potential assembly
+use super::CellGeometry;
+use crate::assembly::common::RlstArray;
+use crate::traits::FunctionSpace;
+use rlst::RlstScalar;
+
+pub trait PotentialIntegrand {
+    //! Integrand
+    /// Scalar type
+    type T: RlstScalar;
+
+    /// Evaluate integrand
+    ///
+    /// # Safety
+    /// This method is unsafe to allow `get_unchecked` to be used
+    unsafe fn evaluate(
+        &self,
+        table: &RlstArray<Self::T, 4>,
+        point_index: usize,
+        eval_index: usize,
+        basis_index: usize,
+        k: &RlstArray<Self::T, 3>,
+        geometry: &impl CellGeometry<T = <Self::T as RlstScalar>::Real>,
+    ) -> Self::T;
+}
+
+pub trait CellAssembler {
+    //! Assembler for the contributions from a cell
+    /// Scalar type
+    type T: RlstScalar;
+
+    /// Assemble contributions into `local_mat`
+    fn assemble(&mut self, local_mat: &mut RlstArray<Self::T, 2>);
+    /// Set the cell
+    fn set_cell(&mut self, cell: usize);
+}
+
+pub trait PotentialAssembly {
+    //! Functions for boundary assembly
+    /// Scalar type
+    type T: RlstScalar;
+
+    /// Assemble into a dense matrix
+    fn assemble_into_dense<Space: FunctionSpace<T = Self::T> + Sync>(
+        &self,
+        output: &mut RlstArray<Self::T, 2>,
+        space: &Space,
+        points: &RlstArray<<Self::T as RlstScalar>::Real, 2>,
+    );
+}


Test patch:
diff --git a/tests/compare_to_bempp_cl.rs b/tests/compare_to_bempp_cl.rs
index d5e3d032..51243c7c 100644
--- a/tests/compare_to_bempp_cl.rs
+++ b/tests/compare_to_bempp_cl.rs
@@ -1,7 +1,7 @@
 use approx::*;
-use bempp::assembly::{boundary::BoundaryAssembler, potential, potential::PotentialAssembler};
+use bempp::assembly::{boundary::BoundaryAssembler, potential::PotentialAssembler};
 use bempp::function::SerialFunctionSpace;
-use bempp::traits::{BoundaryAssembly, FunctionSpace};
+use bempp::traits::{BoundaryAssembly, FunctionSpace, PotentialAssembly};
 use cauchy::c64;
 use ndelement::ciarlet::LagrangeElementFamily;
 use ndelement::types::Continuity;
@@ -237,7 +237,7 @@ fn test_laplace_single_layer_potential_dp0() {
     *points.get_mut([1, 1]).unwrap() = 2.0;
     *points.get_mut([2, 2]).unwrap() = 2.0;
 
-    let a = potential::SingleLayerPotentialAssembler::<f64, _>::new_laplace();
+    let a = PotentialAssembler::<f64, _, _>::new_laplace_single_layer();
     a.assemble_into_dense(&mut matrix, &space, &points);
 
     // Compare to result from bempp-cl
@@ -265,7 +265,7 @@ fn test_helmholtz_single_layer_potential_dp0() {
     *points.get_mut([1, 1]).unwrap() = 2.0;
     *points.get_mut([2, 2]).unwrap() = 2.0;
 
-    let a = potential::SingleLayerPotentialAssembler::<c64, _>::new_helmholtz(3.0);
+    let a = PotentialAssembler::<c64, _, _>::new_helmholtz_single_layer(3.0);
     a.assemble_into_dense(&mut matrix, &space, &points);
 
     // Compare to result from bempp-cl
@@ -293,7 +293,7 @@ fn test_laplace_double_layer_potential_dp0() {
     *points.get_mut([1, 1]).unwrap() = 2.0;
     *points.get_mut([2, 2]).unwrap() = 2.0;
 
-    let a = potential::DoubleLayerPotentialAssembler::<f64, _>::new_laplace();
+    let a = PotentialAssembler::<f64, _, _>::new_laplace_double_layer();
     a.assemble_into_dense(&mut matrix, &space, &points);
 
     // Compare to result from bempp-cl
@@ -321,7 +321,7 @@ fn test_helmholtz_double_layer_potential_dp0() {
     *points.get_mut([1, 1]).unwrap() = 2.0;
     *points.get_mut([2, 2]).unwrap() = 2.0;
 
-    let a = potential::DoubleLayerPotentialAssembler::<c64, _>::new_helmholtz(3.0);
+    let a = PotentialAssembler::<c64, _, _>::new_helmholtz_double_layer(3.0);
     a.assemble_into_dense(&mut matrix, &space, &points);
 
     // Compare to result from bempp-cl