From 53e2c85392b256f22f6267b7ff2819831d57b752 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 16 Apr 2026 00:33:57 -0700 Subject: [PATCH 1/2] Fix global_evaluate parallel hang and shape mismatch (#113) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three related issues fixed: 1. Parallel hang in evaluate_nd: petsc_interpolate had early returns for empty coords that skipped DMLocatePoints — a collective PETSc operation. When some ranks had zero interior points (all OOB), they returned early while others blocked in the collective. Fixed by removing the early returns and ensuring all ranks participate in DMLocatePoints, even with zero points. 2. Shape mismatch in global_evaluate_nd: return arrays were sized to the post-migration particle count rather than the original input count. Pre-allocate with NaN so the shape is always correct regardless of whether points survive the migration round-trip. 3. points_in_domain misclassifying extreme OOB points: the boundary control-point sign heuristic has no distance ceiling, so points at (1e12, 1e12) could be classified as "interior". Added a domain radius check to reject points far beyond the domain extent. Also set ignoreOutsideDomain=True in DMInterpolationSetUp_UW so PETSc silently skips unlocatable points rather than crashing. Underworld development team with AI support from Claude Code --- .../discretisation/discretisation_mesh.py | 18 ++++++- .../function/_dminterp_wrapper.pyx | 33 ++++++++----- src/underworld3/function/_function.pyx | 49 ++++++++++--------- 3 files changed, 63 insertions(+), 37 deletions(-) diff --git a/src/underworld3/discretisation/discretisation_mesh.py b/src/underworld3/discretisation/discretisation_mesh.py index 132ba7b9d..b22f1cffc 100644 --- a/src/underworld3/discretisation/discretisation_mesh.py +++ b/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] diff --git a/src/underworld3/function/_dminterp_wrapper.pyx b/src/underworld3/function/_dminterp_wrapper.pyx index 39d85cc8f..0423c7513 100644 --- a/src/underworld3/function/_dminterp_wrapper.pyx +++ b/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, &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, &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}") diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 98f218e3b..3afd00602 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -484,6 +484,7 @@ 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) @@ -509,13 +510,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 +850,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 +996,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 +1098,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 +1129,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) From cc96d68a7a5a3d465e33b348176f66205690e341 Mon Sep 17 00:00:00 2001 From: lmoresi Date: Thu, 16 Apr 2026 07:54:55 -0700 Subject: [PATCH 2/2] Add geographic return_coords_to_bounds and guard empty-particle case - RegionalGeographicBox now has lateral + radial coordinate clamping via cartesian_to_geographic / geographic_to_cartesian round-trip. Upstream points that escape the domain are clamped to the mesh's lon/lat/depth ranges before evaluation. - Guard empty-particle assignment in global_evaluate_nd: when a rank has zero particles after migration, skip the swarm variable assignment that fails on reshape of empty arrays. Underworld development team with AI support from Claude Code --- src/underworld3/function/_function.pyx | 5 ++- src/underworld3/meshing/geographic.py | 52 ++++++++++++++++++-------- 2 files changed, 40 insertions(+), 17 deletions(-) diff --git a/src/underworld3/function/_function.pyx b/src/underworld3/function/_function.pyx index 3afd00602..3d550ef94 100644 --- a/src/underworld3/function/_function.pyx +++ b/src/underworld3/function/_function.pyx @@ -490,8 +490,9 @@ def global_evaluate_nd( expr, 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[...] diff --git a/src/underworld3/meshing/geographic.py b/src/underworld3/meshing/geographic.py index 77836e94c..70e157ecf 100644 --- a/src/underworld3/meshing/geographic.py +++ b/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, )