Refactored data writing

fenicsxprecice/adapter_core.py

@@ -83,67 +83,19 @@ def determine_function_type(input_obj):
     tag : bool
         0 if input_function is SCALAR and 1 if input_function is VECTOR.
     """
-    if isinstance(input_obj, FunctionSpace):  # scalar-valued functions have rank 0 is FEniCSx
-        if input_obj.num_sub_spaces == 0:
-            return FunctionType.SCALAR
-        elif input_obj.num_sub_spaces == 2:
-            return FunctionType.VECTOR
+    if isinstance(input_obj, FunctionSpace):
+        space = input_obj
     elif isinstance(input_obj, Function):
-        if len(input_obj.x.array.shape) == 1:
-            return FunctionType.SCALAR
-        elif input_obj.x.array.shape[1] > 1:
-            return FunctionType.VECTOR
-        else:
-            raise Exception("Error determining type of given dolfin Function")
+        space = input_obj.function_space
     else:
-        raise Exception("Error determining type of given dolfin FunctionSpace")
-
-
-def convert_fenicsx_to_precice(fenicsx_function, local_ids):
-    """
-    Converts data of type dolfin.Function into Numpy array for all x and y coordinates on the boundary.
-
-    Parameters
-    ----------
-    fenicsx_function : FEniCSx function
-        A FEniCSx function referring to a physical variable in the problem.
-    local_ids: numpy array
-        Array of local indices of vertices on the coupling interface and owned by this rank.
+        raise Exception("Error: determine_function_type must take a Function or FunctionSpace as argument.")
 
-    Returns
-    -------
-    precice_data : array_like
-        Array of FEniCSx function values at each point on the boundary.
-    """
-
-    if not isinstance(fenicsx_function, Function):
-        raise Exception("Cannot handle data type {}".format(type(fenicsx_function)))
-
-    precice_data = []
-    # sampled_data = fenicsx_function.x.array  # that works only for 1st order elements where dofs = grid points
-    # TODO begin dirty fix
-    x_mesh = fenicsx_function.function_space.mesh.geometry.x
-    x_dofs = fenicsx_function.function_space.tabulate_dof_coordinates()
-    mask = []  # where dof coordinate == mesh coordinate
-    for i in range(x_dofs.shape[0]):
-        for j in range(x_mesh.shape[0]):
-            if np.allclose(x_dofs[i, :], x_mesh[j, :], 1e-15):
-                mask.append(i)
-                break
-    sampled_data = fenicsx_function.x.array[mask]
-    # end dirty fix
-
-    if len(local_ids):
-        if fenicsx_function.function_space.num_sub_spaces > 0:  # function space is VectorFunctionSpace
-            for lid in local_ids:
-                precice_data.append(sampled_data[lid, :])
-        else:  # function space is FunctionSpace (scalar)
-            for lid in local_ids:
-                precice_data.append(sampled_data[lid])
+    if space.num_sub_spaces == 0:
+        return FunctionType.SCALAR
+    elif space.num_sub_spaces == 2:
+        return FunctionType.VECTOR
     else:
-        precice_data = np.array([])
-
-    return np.array(precice_data)
+        raise Exception("Error determining type of given dolfin FunctionSpace")
 
 
 def get_fenicsx_vertices(function_space, coupling_subdomain, dims):

fenicsxprecice/fenicsxprecice.py

@@ -7,9 +7,10 @@
 from .config import Config
 import logging
 import precice
-from .adapter_core import FunctionType, determine_function_type, convert_fenicsx_to_precice, get_fenicsx_vertices, CouplingMode, Vertices, VertexType
+from .adapter_core import FunctionType, determine_function_type, get_fenicsx_vertices, CouplingMode, Vertices, VertexType
 from .expression_core import SegregatedRBFInterpolationExpression
 from .solverstate import SolverState
+from .mapping_utils import precompute_eval_vertices
 from dolfinx.fem import Function, FunctionSpace
 import copy
 
@@ -157,7 +158,7 @@ def read_data(self):
 
         return copy.deepcopy(read_data)
 
-    def write_data(self, write_function):
+    def write_data(self, write_function: Function):
         """
         Writes data to preCICE. Depending on the dimensions of the simulation (2D-3D Coupling, 2D-2D coupling or
         Scalar/Vector write function) write_data is first converted into a format needed for preCICE.
@@ -172,18 +173,21 @@ def write_data(self, write_function):
                 CouplingMode.BI_DIRECTIONAL_COUPLING)
 
         w_func = write_function
+        write_function_type = determine_function_type(write_function)
 
         # Check that the function provided lives on the same function space provided during initialization
-        assert (self._write_function_type == determine_function_type(w_func))
-        # TODO this raises AssertionError, not sure why. I just commented it out, still works...
-        # assert (write_function.function_space == self._write_function_space)
+        assert (self._write_function_type == write_function_type)
+        # TODO: fails for weird reasons.
+        #assert (write_function.function_space == self._write_function_space)
 
         write_data_id = self._interface.get_data_id(self._config.get_write_data_name(),
                                                     self._interface.get_mesh_id(self._config.get_coupling_mesh_name()))
 
-        write_function_type = determine_function_type(write_function)
         assert (write_function_type in list(FunctionType))
-        write_data = convert_fenicsx_to_precice(write_function, self._fenicsx_vertices.get_ids())
+
+        # Eval the function on the vertices
+        write_data = w_func.eval(self._fenicsx_vertices_eval, self._fenicsx_cells_eval)
+
         if write_function_type is FunctionType.SCALAR:
             assert (write_function.function_space.num_sub_spaces == 0)
             write_data = np.squeeze(write_data)  # TODO dirty solution
@@ -289,6 +293,12 @@ def initialize(self, coupling_subdomain, read_function_space=None, write_object=
         self._fenicsx_vertices.set_ids(ids)
         self._fenicsx_vertices.set_coordinates(coords)
 
+        # Set up helpers to write data to preCICE:
+        if self._coupling_type is CouplingMode.UNI_DIRECTIONAL_WRITE_COUPLING or \
+                self._coupling_type is CouplingMode.BI_DIRECTIONAL_COUPLING:
+            self._fenicsx_vertices_eval, self._fenicsx_cells_eval = precompute_eval_vertices(
+                self._fenicsx_vertices.get_coordinates(), write_function_space.mesh)
+
         # Set up mesh in preCICE
         self._precice_vertex_ids = self._interface.set_mesh_vertices(self._interface.get_mesh_id(
             self._config.get_coupling_mesh_name()), self._fenicsx_vertices.get_coordinates())

fenicsxprecice/mapping_utils.py

@@ -0,0 +1,27 @@
+"""
+This module consists of helper functions used in the adapter related to data mapping. Names of the functions are self explanatory
+"""
+
+from dolfinx.fem import FunctionSpace, Function
+from dolfinx.mesh import Mesh
+from dolfinx.geometry import BoundingBoxTree, compute_colliding_cells, compute_collisions
+import numpy as np
+
+
+def precompute_eval_vertices(precice_vertices: np.ndarray, mesh: Mesh):
+    # Pre-processing: pad with zeros if 2D vector
+    n_vertices, dim = precice_vertices.shape
+    if dim == 2:
+        precice_vertices = np.append(
+            precice_vertices, np.zeros((n_vertices, 1)), axis=1)
+
+    assert precice_vertices.shape[1] == 3
+
+    tree = BoundingBoxTree(mesh, mesh.geometry.dim)
+    cell_candidates = compute_collisions(tree, precice_vertices)
+    cell = compute_colliding_cells(mesh, cell_candidates, precice_vertices)
+
+    # Adjacency list: convert to array and take the first match for each item.
+    # Use offsets (but the last one) to compute indices in the array
+    list_of_cells = cell.array[cell.offsets[:-1]]
+    return precice_vertices, list_of_cells

tests/integration/test_fenicsxprecice.py

@@ -158,13 +158,13 @@ def right_boundary(x): return abs(x[0] - 1.0) < 10**-14
         precice._interface = Interface(None, None, None, None)
         precice.initialize(right_boundary, self.scalar_V, self.scalar_function)
         values = np.array([self.scalar_function.eval([x, y, 0], 0)[0]
-                          for x, y in zip(self.vertices_x, self.vertices_y)])
+                           for x, y in zip(self.vertices_x, self.vertices_y)])
         data = {(x, y): v for x, y, v in zip(self.vertices_x, self.vertices_y, values)}
         scalar_coupling_expr = precice.create_coupling_expression()
         precice.update_coupling_expression(scalar_coupling_expr, data)
 
         expr_samples = np.array([scalar_coupling_expr.eval([x, y, 0], 0)
-                                for x, y in zip(self.samplepts_x, self.samplepts_y)])
+                                 for x, y in zip(self.samplepts_x, self.samplepts_y)])
         func_samples = np.array([self.scalar_function.eval([x, y, 0], 0)
                                  for x, y in zip(self.samplepts_x, self.samplepts_y)])
 
@@ -201,7 +201,7 @@ def right_boundary(x): return abs(x[0] - 1.0) < 10**-14
         precice.update_coupling_expression(vector_coupling_expr, data)
 
         expr_samples = np.array([vector_coupling_expr.eval([x, y, 0], 0)
-                                for x, y in zip(self.samplepts_x, self.samplepts_y)])
+                                 for x, y in zip(self.samplepts_x, self.samplepts_y)])
         func_samples = np.array([self.vector_function.eval([x, y, 0], 0)
                                  for x, y in zip(self.samplepts_x, self.samplepts_y)])
 

tests/integration/test_write_read.py

@@ -85,7 +85,7 @@ def test_vector_write(self):
         """
         from precice import Interface
         import fenicsxprecice
-        from fenicsxprecice.adapter_core import VertexType, Vertices, convert_fenicsx_to_precice
+        from fenicsxprecice.adapter_core import VertexType, Vertices
 
         Interface.write_block_vector_data = MagicMock()
         Interface.get_dimensions = MagicMock(return_value=self.dimension)

tests/unit/test_mapping_utils.py

@@ -0,0 +1,46 @@
+import unittest
+from unittest.mock import MagicMock
+from unittest import TestCase
+import numpy as np
+from mpi4py import MPI
+from dolfinx.mesh import create_unit_square
+from dolfinx.fem import Function, FunctionSpace, VectorFunctionSpace
+
+
+class TestAdapterCore(TestCase):
+    def test_precompute_eval_vertices(self):
+        """
+        Test cell collision computation for function evaluation on vertices
+        """
+        # TODO: handle preCIE using N*2 arrays and not N*3 in 2D simulations
+        from fenicsxprecice.mapping_utils import precompute_eval_vertices
+
+        mesh = create_unit_square(MPI.COMM_WORLD, 2, 2)  # create dummy mesh
+        precice_vertices = np.array([[0.5, 0.5, 0.0],
+                                     [0.2, 0.2, 0.0],
+                                     [1.0, 1.0, 0.0]])
+
+        # eval_pos is a copy of precice_vertices but with padded 0s in the z-dim
+        eval_pos, cells = precompute_eval_vertices(precice_vertices, mesh)
+
+        # scalar valued
+        V = FunctionSpace(mesh, ('P', 2))  # Create function space using mesh
+
+        fenicsx_function = Function(V)
+        fenicsx_function.interpolate(lambda x: x[0] + x[1] * x[1])
+
+        expected_data = np.array([0.75, 0.24, 2.0]).reshape((3, 1))
+
+        data = fenicsx_function.eval(eval_pos, cells)
+        np.testing.assert_almost_equal(data, expected_data)
+
+        # Vector valued
+        V = VectorFunctionSpace(mesh, ('P', 2))
+
+        fenicsx_function = Function(V)
+        fenicsx_function.interpolate(lambda x: (x[0] + x[1] * x[1], x[0]))
+
+        expected_data = np.array([[0.75, 0.50], [0.24, 0.20], [2.0, 1.0]])
+
+        data = fenicsx_function.eval(eval_pos, cells)
+        np.testing.assert_almost_equal(data, expected_data)