Merge pull request #708 from pangeo-forge/multidim-rechunk-fix

Minimal multidim rechunk fix

pangeo_forge_recipes/rechunking.py

@@ -146,7 +146,8 @@ def _invert_meshgrid(*arrays):
         selectors[n] = tuple(selectors[n])
     xi = [a[s] for a, s in zip(arrays, selectors)]
     assert all(
-        np.equal(actual, expected).all() for actual, expected in zip(arrays, np.meshgrid(*xi))
+        np.equal(actual, expected.squeeze()).all()
+        for actual, expected in zip(arrays, np.meshgrid(*xi))
     )
     return xi
 

tests/test_rechunking.py

@@ -26,10 +26,11 @@
     [(5, "1D"), (10, "2D")],
 )
 @pytest.mark.parametrize("time_chunks", [1, 2, 5])
-def test_split_and_combine_fragments_with_merge_dim(nt_dayparam, time_chunks):
+@pytest.mark.parametrize("other_chunks", [{}, {"lat": 5}, {"lat": 5, "lon": 5}])
+def test_split_and_combine_fragments_with_merge_dim(nt_dayparam, time_chunks, other_chunks):
     """Test if sub-fragments split from datasets with merge dims can be combined with each other."""
 
-    target_chunks = {"time": time_chunks}
+    target_chunks = {"time": time_chunks, **other_chunks}
     nt, dayparam = nt_dayparam
     ds = make_ds(nt=nt)
     dsets, _, _ = split_up_files_by_variable_and_day(ds, dayparam)