Fixes regrid2 mapping output to input ordering (#653)

tests/test_regrid.py

@@ -372,6 +372,30 @@ def setup(self):
  }
  )
 
+ # source ordering is time, height, lat, lon
+ @pytest.mark.parametrize(
+ "ordering",
+ [
+ ("time", "lat", "height", "lon"),
+ ("lat", "lon", "height", "time"),
+ ("lat", "height", "time", "lon"),
+ # no change in ordering
+ ("time", "height", "lat", "lon"),
+ ],
+ )
+ def test_ordering(self, ordering):
+ ds = self.coarse_4d_ds
+
+ output_grid = grid.create_gaussian_grid(4)
+
+ ds["ts"] = ds.ts.transpose(*ordering)
+
+ regridder = regrid2.Regrid2Regridder(ds, output_grid)
+
+ output_ds = regridder.horizontal("ts", ds)
+
+ assert ds.ts.dims == output_ds.ts.dims
+
  def test_vertical_placeholder(self):
  ds = fixtures.generate_dataset(
  decode_times=True, cf_compliant=False, has_bounds=True

xcdat/regridder/regrid2.py

@@ -150,6 +150,7 @@ def _regrid(
  other_sizes = list(other_dims.values())
 
  data_shape = [y_length * x_length] + other_sizes
+
  # output data is always float32 in original code
  output_data = np.zeros(data_shape, dtype=np.float32)
  output_mask = np.ones(data_shape, dtype=np.float32)
@@ -212,7 +213,11 @@ def _regrid(
 
  output_data = output_data.reshape(output_data_shape)
 
- output_order = [x + 2 for x in range(input_data_var.ndim - 2)] + [0, 1]
+ # temp dimensional ordering
+ temp_dims = [y_name, x_name] + list(other_dims.keys())
+
+ # map temp ordering to input ordering
+ output_order = [temp_dims.index(x) for x in input_data_var.dims]
 
  output_data = output_data.transpose(output_order)