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..3d550ef94 100644 --- a/src/underworld3/function/_function.pyx +++ b/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) 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, )