Fix global_evaluate parallel hang, shape mismatch, and OOB classification

src/underworld3/discretisation/discretisation_mesh.py

@@ -2568,12 +2568,20 @@ def _mark_local_boundary_faces_inside_and_out(self):
                 control_points_list.append(inside_control_point)
                 control_point_sign_list.append(1)
 
-        control_point_kdtree = uw.kdtree.KDTree(numpy.array(control_points_list))
+        control_points_array = numpy.array(control_points_list)
+        control_point_kdtree = uw.kdtree.KDTree(control_points_array)
         control_point_sign = numpy.array(control_point_sign_list)
 
         self.boundary_face_control_points_kdtree = control_point_kdtree
         self.boundary_face_control_points_sign = control_point_sign
 
+        # Domain bounding radius (squared): distance from centroid to farthest
+        # control point. Points beyond this distance from their nearest control
+        # point cannot be inside the domain.
+        domain_centroid = control_points_array.mean(axis=0)
+        radii_sq = numpy.sum((control_points_array - domain_centroid) ** 2, axis=1)
+        self._domain_radius_squared = float(radii_sq.max())
+
         return
 
     def points_in_domain(self, points, strict_validation=True):
@@ -2611,10 +2619,16 @@ def points_in_domain(self, points, strict_validation=True):
         dist2, closest_control_points_ext = self.boundary_face_control_points_kdtree.query(
             model_points, k=1, sqr_dists=True
         )
+        dist2 = numpy.asarray(dist2).ravel()  # kd-tree returns (n,1) for k=1
         in_or_not = self.boundary_face_control_points_sign[closest_control_points_ext] > 0
 
-        ## This choice of distance needs some more thought
+        # Points very far from the nearest boundary face are definitely exterior.
+        # The sign heuristic only works for points within the domain's neighbourhood;
+        # beyond that, "nearest control point" is arbitrary.
+        far_from_domain = dist2 > self._domain_radius_squared
+        in_or_not[far_from_domain] = False
 
+        # Points close to the boundary need the expensive cell-location check
         near_boundary = numpy.where(dist2 < 2 * max_radius**2)[0]
         near_boundary_points = model_points[near_boundary]
 

src/underworld3/function/_dminterp_wrapper.pyx

@@ -121,21 +121,30 @@ cdef class CachedDMInterpolationInfo:
             DMInterpolationDestroy(&self._ipInfo)
             raise RuntimeError(f"DMInterpolationSetDof failed with error {ierr}")
 
-        # Add interpolation points - use contiguous array's data pointer
-        cdef double[:, ::1] coords_view = np.ascontiguousarray(self.coords)
-        ierr = DMInterpolationAddPoints(self._ipInfo, n_points, &coords_view[0, 0])
-        if ierr != 0:
-            DMInterpolationDestroy(&self._ipInfo)
-            raise RuntimeError(f"DMInterpolationAddPoints failed with error {ierr}")
-
-        # Set up with cell hints
-        # Extract PETSc DM from mesh
+        # Declare typed memoryviews at function scope (Cython requirement)
+        cdef double[:, ::1] coords_view
+        cdef long[::1] cells_view
         cdef DM dm_obj = mesh.dm
         cdef PetscDM dm = dm_obj.dm
 
-        # Extract cell hints as size_t array
-        cdef long[::1] cells_view = np.ascontiguousarray(self.cells)
-        ierr = DMInterpolationSetUp_UW(self._ipInfo, dm, 0, 0, <size_t*> &cells_view[0])
+        # Add interpolation points (guard against empty arrays)
+        if n_points > 0:
+            coords_view = np.ascontiguousarray(self.coords)
+            ierr = DMInterpolationAddPoints(self._ipInfo, n_points, &coords_view[0, 0])
+            if ierr != 0:
+                DMInterpolationDestroy(&self._ipInfo)
+                raise RuntimeError(f"DMInterpolationAddPoints failed with error {ierr}")
+
+        # Set up — calls DMLocatePoints which is COLLECTIVE on the mesh DM.
+        # All ranks must call this, even with zero local points.
+        # ignoreOutsideDomain=1: PETSc silently skips points it cannot locate
+        # rather than crashing. This is essential — points_in_domain() uses a
+        # kd-tree heuristic that can misclassify distant points as interior.
+        if n_points > 0:
+            cells_view = np.ascontiguousarray(self.cells)
+            ierr = DMInterpolationSetUp_UW(self._ipInfo, dm, 0, 1, <size_t*> &cells_view[0])
+        else:
+            ierr = DMInterpolationSetUp_UW(self._ipInfo, dm, 0, 1, NULL)
         if ierr != 0:
             DMInterpolationDestroy(&self._ipInfo)
             raise RuntimeError(f"DMInterpolationSetUp_UW failed with error {ierr}")

src/underworld3/function/_function.pyx

@@ -484,13 +484,15 @@ def global_evaluate_nd(   expr,
 
     index = original_index.array[:,0,0]
     ranks = original_rank.array[:,0,0]
+    n_input_points = coords_array.shape[0]
 
     evaluation_swarm.migrate(remove_sent_points=True, delete_lost_points=False)
     local_coords = evaluation_swarm._particle_coordinates.array[...].reshape(-1,evaluation_swarm.dim)
     values, extrapolated = evaluate_nd(expr, local_coords, rbf=rbf, evalf=evalf, verbose=verbose, check_extrapolated=True,)
 
-    data_container.array[...] = values[...]
-    is_extrapolated.array[:,0,0] = extrapolated[:]
+    if local_coords.shape[0] > 0:
+        data_container.array[...] = values[...]
+        is_extrapolated.array[:,0,0] = extrapolated[:]
 
     # set rank to old values and migrate back
     evaluation_swarm._rank_var.array[...] = original_rank.array[...]
@@ -509,13 +511,17 @@ def global_evaluate_nd(   expr,
         if hasattr(var, "_canonical_data"):
             var._canonical_data = None
 
-    index = original_index.array[:,0,0]
-
-    return_value = np.empty_like(data_container.array[...])
-    return_value[index,:,:] =  data_container.array[:,:,:]
+    # Pre-allocate with NaN so the shape is always correct. If any points
+    # are lost during the migration round-trip, they remain NaN rather than
+    # causing a shape mismatch or returning uninitialised data.
+    return_value = np.full((n_input_points,) + expr_shape, np.nan, dtype=np.double)
+    return_mask = np.full((n_input_points, 1, 1), True, dtype=bool)
 
-    return_mask = np.empty_like(is_extrapolated.array[...], dtype=bool)
-    return_mask[index] =  is_extrapolated.array[:]
+    n_returned = original_index.array.shape[0]
+    if n_returned > 0:
+        index = original_index.array[:, 0, 0].astype(int)
+        return_value[index, :, :] = data_container.array[:, :, :]
+        return_mask[index] = is_extrapolated.array[:]
 
     if not check_extrapolated:
         return return_value
@@ -845,6 +851,12 @@ def evaluate_nd(   expr,
                             )
 
     else:
+        # CRITICAL: update_lvec() calls dm.globalToLocal() which is COLLECTIVE.
+        # It MUST be called by ALL ranks before any rank enters petsc_interpolate,
+        # because ranks with zero interior points would skip petsc_interpolate
+        # (and its internal update_lvec call), deadlocking the ranks that do enter.
+        mesh.update_lvec()
+
         in_or_not = mesh.points_in_domain(coords_array, strict_validation=False)
         evaluation_interior = petsc_interpolate( expr,
                                     coords_array[in_or_not],
@@ -985,18 +997,11 @@ def petsc_interpolate(   expr,
     if other_arguments:
         raise RuntimeError("`other_arguments` functionality not yet implemented.")
 
-    # Early return for empty coordinate arrays (SECOND CHECK - top-level function)
-    # CRITICAL: Avoid lambdify errors with LaTeX variable names when coords is empty
-    # This handles cases where empty arrays pass through from evaluate_nd
-    if len(coords) == 0:
-        # Determine output shape based on expression type
-        try:
-            expr_shape = expr.shape
-            # Return empty array with correct shape: (0, rows, cols)
-            return np.empty([0] + list(expr_shape), dtype=np.double)
-        except AttributeError:
-            # Scalar expression - return (0,) shaped array
-            return np.empty([0], dtype=np.double)
+    # NOTE: Do NOT early-return for empty coords here. petsc_interpolate
+    # calls DMLocatePoints which is COLLECTIVE on the mesh DM communicator.
+    # If some ranks skip it (empty coords) while others enter it, MPI deadlocks.
+    # Empty coords are handled inside interpolate_vars_on_mesh after the
+    # collective operations complete.
 
     ## Substitute any UWExpressions for their values before calculation
     ## NOTE: We use _unwrap_expressions directly (not fn_substitute_expressions) to avoid
@@ -1094,11 +1099,9 @@ def petsc_interpolate(   expr,
         # Make coords contiguous for caching and C access
         coords = np.ascontiguousarray(coords)
 
-        # Early return for empty coordinate arrays
-        # CRITICAL: Avoid DMInterpolation setup with zero points
-        if len(coords) == 0:
-            # Return empty array with correct shape: (0, dofcount)
-            return np.empty([0, dofcount], dtype=np.double)
+        # NOTE: No early return for empty coords here. DMLocatePoints
+        # (inside DMInterpolationSetUp_UW) is COLLECTIVE on the mesh DM
+        # communicator. All ranks must participate, even with zero points.
 
         # === DMInterpolation CACHING ===
         # Declare variables at function scope (Cython requirement)
@@ -1127,6 +1130,7 @@ def petsc_interpolate(   expr,
             cells = mesh.get_closest_cells(coords)
 
             # Create and set up DMInterpolation structure (EXPENSIVE)
+            # This calls DMLocatePoints which is COLLECTIVE — all ranks must enter.
             try:
                 # coords is already np.ndarray type (function signature ensures this)
                 cached_info.create_structure(mesh, coords, cells, dofcount)

src/underworld3/meshing/geographic.py

@@ -329,21 +329,6 @@ class boundaries(Enum):
         gmsh.write(uw_filename)
         gmsh.finalize()
 
-    ## This needs a side-boundary capture routine as well
-
-    def sphere_return_coords_to_bounds(coords):
-        Rsq = coords[:, 0] ** 2 + coords[:, 1] ** 2 + coords[:, 2] ** 2
-
-        outside = Rsq > radiusOuter**2
-        inside = Rsq < radiusInner**2
-
-        ## Note these numbers should not be hard-wired
-
-        coords[outside, :] *= 0.99 * radiusOuter / np.sqrt(Rsq[outside].reshape(-1, 1))
-        coords[inside, :] *= 1.01 * radiusInner / np.sqrt(Rsq[inside].reshape(-1, 1))
-
-        return coords
-
     def spherical_mesh_refinement_callback(dm):
         r_o = radiusOuter
         r_i = radiusInner
@@ -789,6 +774,42 @@ class boundaries(Enum):
         gmsh.write(uw_filename)
         gmsh.finalize()
 
+    def geographic_return_coords_to_bounds(coords):
+        """Clamp Cartesian coordinates to the geographic domain bounds.
+
+        Converts to geographic (lon, lat, depth), clamps each to the
+        mesh's known ranges, and converts back to Cartesian. Small
+        overshoot due to topography or mesh curvature is handled
+        gracefully by the interior/exterior split in evaluate_nd.
+
+        Coords must be in the mesh's internal coordinate system
+        (nondimensional if scaling is active, km otherwise).
+        """
+        from underworld3.coordinates import cartesian_to_geographic
+
+        # Work with raw numpy float arrays — coords may be UnitAwareArray
+        raw = np.asarray(coords, dtype=np.float64)
+
+        lon, lat, depth = cartesian_to_geographic(
+            raw[:, 0], raw[:, 1], raw[:, 2], float(a), float(b)
+        )
+
+        # Ensure plain float arrays for clip
+        lon = np.asarray(lon, dtype=np.float64)
+        lat = np.asarray(lat, dtype=np.float64)
+        depth = np.asarray(depth, dtype=np.float64)
+
+        np.clip(lon, lon_min, lon_max, out=lon)
+        np.clip(lat, lat_min, lat_max, out=lat)
+        np.clip(depth, float(depth_min) * 1.01, float(depth_max) * 0.99, out=depth)
+
+        x, y, z = geographic_to_cartesian(lon, lat, depth, a, b)
+        coords[:, 0] = x
+        coords[:, 1] = y
+        coords[:, 2] = z
+
+        return coords
+
     # Load mesh on all ranks
     new_mesh = Mesh(
         uw_filename,
@@ -802,6 +823,7 @@ class boundaries(Enum):
         refinement=refinement,
         coarsening=coarsening,
         coordinate_system_type=CoordinateSystemType.GEOGRAPHIC,
+        return_coords_to_bounds=geographic_return_coords_to_bounds,
         verbose=verbose,
     )