Creating an "opaque" interpolant object of a FEM solution

I need to create a “callable” object from the solution of a FEM problem so I can use it as an initial condition for a coupled problem with a perturbed mesh, i.e. I can’t just copy over the current values at the domain DOFs.

I tried to create such a callable object and it works when -np is 1, but deadlocks when running with more than one process.

class InterpolatedSolution:
    def __init__(self, fem_solution, domain):
        self.fem_solution = fem_solution
        self.bb_tree = geometry.bb_tree(domain, domain.topology.dim)
        self.domain = domain

    def __calc_solution_at(self, x):
        pt = np.array([x]).reshape(1, 3)
        
        # Choose one of the cells that contains the point
        cell_candidates = geometry.compute_collisions_points(self.bb_tree, pt)
        colliding_cells = geometry.compute_colliding_cells(
            self.domain, cell_candidates, pt
        )
        
        cells = []
        if len(colliding_cells.links(0)) > 0:
            cells.append(colliding_cells.links(0)[0])
        
        result = 0.0
        if len(cells) == 1:
            result = float(np.squeeze(self.fem_solution.eval(pt, cells)))

        comm.barrier()  # memory fence to make sure we wait on the actual thread that does the calculation
        # from here on, result should be either 0.0 or the actual value on all threads
        
        output = comm.allreduce(result, op=MPI.SUM) 
        return output

    def __call__(self, x):
        x = np.array(x)
        if x.shape == (3,):
            return self.__calc_solution_at(x)

        return np.apply_along_axis(self.__calc_solution_at, axis=0, arr=x)

I’m sure that there is an easier way to do this, but I’m not used to MPI/fork-style parallelism yet :smiley: