dev

davidhassell · davidhassell · commit ebd47b4a3e0f · 2023-09-20T13:44:38.000+01:00
diff --git a/cf/data/dask_regrid.py b/cf/data/dask_regrid.py
@@ -537,8 +537,8 @@ def _regrid(
     a = weights.dot(a)
 
     if dst_mask is not None:
-        a = np.ma.array(a, mask=dst_mask)
-    #        a[dst_mask] = np.ma.masked
+        a = np.ma.array(a)
+        a[dst_mask] = np.ma.masked
 
     return a, src_mask, dst_mask, weights
 
diff --git a/cf/data/data.py b/cf/data/data.py
@@ -3754,10 +3754,33 @@ def _regrid(
             dx = dx.rechunk(chunks)
 
         # Define the regridded chunksizes
-        regridded_chunks = tuple(
-            (regridded_sizes[i],) if i in regridded_sizes else c
-            for i, c in enumerate(dx.chunks)
-        )
+        regridded_chunks = []  # The 'chunks' parameter to `map_blocks`
+        new_axis = []  # The 'new_axis' parameter to `map_blocks`
+        n = 0
+        for i, c in enumerate(dx.chunks):
+            if i in regridded_sizes:
+                sizes = regridded_sizes[i]
+                n_sizes = len(sizes)
+                regridded_chunks.extend(sizes)
+                if n_sizes > 1:
+                    new_axis.extend(range(n + 1, n + n_sizes))
+                    n += n_sizes - 1
+            else:
+                regridded_chunks.extend(c)
+
+            n += 1
+
+        if new_axis:
+            # Update the axis identifiers.
+            #
+            # This is necessary when regridding changes the number of
+            # data dimensions (e.g. as happens when regridding a mesh
+            # topology axis to separate lat and lon axes).
+            axes = list(self._axes)
+            for i in new_axis:
+                axes.insert(i, new_axis_identifier(tuple(axes)))
+
+            self._axes = tuple(axes)
 
         # Set the output data type
         if method in ("nearest_dtos", "nearest_stod"):
@@ -3790,6 +3813,7 @@ def _regrid(
             weights_dst_mask=weights_dst_mask,
             ref_src_mask=src_mask,
             chunks=regridded_chunks,
+            new_axis=new_axis,
             meta=np.array((), dtype=dst_dtype),
         )
 
diff --git a/cf/regrid/regrid.py b/cf/regrid/regrid.py
