diff --git a/pangeo_forge_recipes/rechunking.py b/pangeo_forge_recipes/rechunking.py index bab9fcc9..26fbf3d6 100644 --- a/pangeo_forge_recipes/rechunking.py +++ b/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 diff --git a/tests/test_rechunking.py b/tests/test_rechunking.py index f8caab9c..d4f644b6 100644 --- a/tests/test_rechunking.py +++ b/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)