fix: swarm empty-rank save deadlock (MPIO dtype) and estimate_dt reshape crash

src/underworld3/swarm.py

@@ -343,6 +343,7 @@ def __init__(
             raise TypeError(
                 f"Provided dtype={dtype} is not supported. Supported types are 'int' and 'float'."
             )
+        self._petsc_dtype = petsc_type
 
         if _register:
             # Check if swarm is already populated - PETSc doesn't allow registering
@@ -450,7 +451,9 @@ def _create_canonical_data_array(self, initial_data=None):
 
             # Handle case where unpack returns None (swarm not initialized)
             if initial_data is None:
-                initial_data = np.zeros((0, self.num_components))
+                initial_data = np.zeros(
+                    (0, self.num_components), dtype=self._petsc_dtype
+                )
 
         # Create NDArray_With_Callback for flat data
         array_obj = uw.utilities.NDArray_With_Callback(
@@ -1428,15 +1431,19 @@ def unpack_raw_data_from_petsc(self, squeeze=True, sync=None):
         # Check if swarm has any particles before accessing field
         swarm_size = self.swarm.local_size
         if swarm_size <= 0:
-            # Swarm not populated yet, return empty array
-            return np.zeros((0, self.num_components))
+            # Swarm not populated yet, return empty array. Keep the field's
+            # PETSc dtype so that an empty rank's array agrees with a
+            # non-empty rank's (a float64 default here made the collective
+            # parallel-HDF5 create_dataset in ``save`` see different dtypes
+            # across ranks and deadlock on close for ``int`` variables).
+            return np.zeros((0, self.num_components), dtype=self._petsc_dtype)
 
         # Direct PETSc field access without context manager
         field_data = self.swarm.dm.getField(self.clean_name)
         if field_data is None:
             # Field not properly initialized, restore and return empty array
             self.swarm.dm.restoreField(self.clean_name)
-            return np.zeros((0, self.num_components))
+            return np.zeros((0, self.num_components), dtype=self._petsc_dtype)
 
         petsc_data = field_data.reshape((-1, self.num_components))
 
@@ -5083,7 +5090,15 @@ def estimate_dt(self, V_fn):
         # silently disabling advection's step_limit substepping (BF-16).
         vel = np.asarray(vel)
         if vel.ndim == 3:
-            vel = vel.reshape(vel.shape[0], -1)
+            # Guard against empty ranks: an array of size 0 cannot be
+            # reshaped with a `-1` axis (NumPy cannot infer the implied
+            # dimension from zero elements) — e.g. (0, 1, dim) -> (0, -1)
+            # raises ValueError. A zero-particle rank legitimately has no
+            # velocities and contributes 0 to the global max below.
+            if vel.size == 0:
+                vel = np.zeros((0, vel.shape[2]) if vel.ndim >= 3 else (0,))
+            else:
+                vel = vel.reshape(vel.shape[0], -1)
 
         try:
             magvel_squared = vel[:, 0] ** 2 + vel[:, 1] ** 2