Repository: xCDAT/xcdat

Issue:
[Bug]: Regrid2 breaks in `_preserve_bounds` because `_regrid` incorrectly transposes `output_data` shape, which doesn't align with original dim order used by bounds
### What happened?

There is an issue where transposing the shape of data causes misalignment with the original dimension order. When we attempt to preserve bounds, the original dimension order is still being used which does not align with the transposed dimension order.


### Example

Let's say we have `ds_a` with an original dimension of shape of `(lat: 1, lev: 2, lon: 3)`. 

After transposing the data, `output_data` become `shape=(2, 1, 3)`.

https://github.com/xCDAT/xcdat/blob/d0e33aaea04ffeca0301cd7773e4f981c7957653/xcdat/regridder/regrid2.py#L164-L172

We recreate the `xr.DataArray` (`output_da`) using `output_data` and `dims`. As mentioned above, the shape of `output_data` is now transposed, but it does not align with the original `dims`. This results in `output_da` having a new dimension shape of `(lat: 2, lev: 1, lon: 3)`
https://github.com/xCDAT/xcdat/blob/d0e33aaea04ffeca0301cd7773e4f981c7957653/xcdat/regridder/regrid2.py#L184-L197

When we attempt to preserve `lat_bnds` in `_preserve_bounds()`, the original `lat_bnds` still has has a shape of `(lat: 1, lev: 2, lon: 3)`, which does not align with `(lat: 2, lev: 1, lon: 3)`.

https://github.com/xCDAT/xcdat/blob/d0e33aaea04ffeca0301cd7773e4f981c7957653/xcdat/regridder/base.py#L48-L58

This results in `ValueError: cannot reindex or align along dimension 'lat' because of conflicting dimension sizes: {1, 2} (note: an index is found along that dimension with size=1)`.


### What did you expect to happen? Are there are possible answers you came across?

Maybe `output_data` shape should align with the original dim order or we should also transpose the shape of the bounds?

### Minimal Complete Verifiable Example (MVCE)

```python
import numpy as np
import xarray as xr
import xcdat as xc

# 1. Create two datasets with dimensions (lat=1, lev=2, lon=3)
ds_a = xr.Dataset(
    {
        "lat": xr.DataArray(
            name="lat", dims="lat", data=[0], attrs={"axis": "Y", "bounds": "lat_bnds"}
        ),
        "lev": xr.DataArray(name="lev", dims="lev", data=[0, 1], attrs={"axis": "Z"}),
        "lon": xr.DataArray(
            name="lon", dims="lon", data=[0, 1, 2], attrs={"axis": "X"}
        ),
    }
)
ds_a["T"] = xr.DataArray(name="T", dims=["lat", "lev", "lon"], data=np.zeros((1, 2, 3)))
ds_a = ds_a.bounds.add_missing_bounds(axes=["X", "Y", "Z"])

# Manually add lat_bnds because xCDAT does not support adding bounds for
# singleton coords (yet)
ds_a["lat_bnds"] = xr.DataArray(name="lat_bnds", dims=["lat", "bnds"], data=[[-1, 1]])


ds_b = ds_a.copy()

# 2. Regrid dataset b to a using regrid2, bilinear
output_grid = ds_a.regridder.grid
ds_b_regrid = ds_b.regridder.horizontal(
    "T", output_grid, tool="regrid2", method="bilinear"
)
# ValueError: cannot reindex or align along dimension 'lat' because of conflicting dimension sizes: {1, 2} (note: an index is found along that dimension with size=1)

# 3. Breaks because `_regrid` incorrectly reorders dimensions
# Resulting in the `output_data` shape to become (lat=2, lev=1, lon=3).
˘
```


### Relevant log output

```python
ValueError                                Traceback (most recent call last)
File /global/u2/v/vo13/E3SM-Project/e3sm_diags/auxiliary_tools/cdat_regression_testing/657-meridional-mean-2d/657_meridional_mean_2d_run_script.py:19
     16 ds_b = ds_a.copy()
     18 output_grid = ds_a.regridder.grid
---> 19 ds_b_regrid = ds_b.regridder.horizontal(
     20     "T", output_grid, tool="regrid2", method="bilinear"
     21 )

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xcdat/regridder/accessor.py:324, in RegridderAccessor.horizontal(self, data_var, output_grid, tool, **options)
    322 input_grid = _get_input_grid(self._ds, data_var, ["X", "Y"])
    323 regridder = regrid_tool(input_grid, output_grid, **options)
--> 324 output_ds = regridder.horizontal(data_var, self._ds)
    326 return output_ds

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xcdat/regridder/regrid2.py:85, in Regrid2Regridder.horizontal(self, data_var, ds)
     79     input_data_var = input_data_var.where(src_mask != 0.0, masked_value)
     81 output_data = _regrid(
     82     input_data_var, src_lat_bnds, src_lon_bnds, dst_lat_bnds, dst_lon_bnds
     83 )
---> 85 output_ds = _build_dataset(
     86     ds,
     87     data_var,
     88     output_data,
     89     dst_lat_bnds,
     90     dst_lon_bnds,
     91     self._input_grid,
     92     self._output_grid,
     93 )
     95 return output_ds

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xcdat/regridder/regrid2.py:206, in _build_dataset(ds, data_var, output_data, dst_lat_bnds, dst_lon_bnds, input_grid, output_grid)
    199 output_data_vars[data_var] = output_da
    201 output_ds = xr.Dataset(
    202     output_data_vars,
    203     attrs=input_grid.attrs.copy(),
    204 )
--> 206 output_ds = _preserve_bounds(ds, output_grid, output_ds, ["X", "Y"])
    208 return output_ds

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xcdat/regridder/base.py:56, in _preserve_bounds(input_ds, output_grid, output_ds, drop_axis)
     54         else:
     55             if bnds.name not in output_ds:
---> 56                 output_ds[bnds.name] = bnds.copy()
     58 return output_ds

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xarray/core/dataset.py:1603, in Dataset.__setitem__(self, key, value)
   1598     if isinstance(value, Dataset):
   1599         raise TypeError(
   1600             "Cannot assign a Dataset to a single key - only a DataArray or Variable "
   1601             "object can be stored under a single key."
   1602         )
-> 1603     self.update({key: value})
   1605 elif utils.iterable_of_hashable(key):
   1606     keylist = list(key)

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xarray/core/dataset.py:5617, in Dataset.update(self, other)
   5581 def update(self, other: CoercibleMapping) -> Self:
   5582     """Update this dataset's variables with those from another dataset.
   5583 
   5584     Just like :py:meth:`dict.update` this is a in-place operation.
   (...)
   5615     Dataset.merge
   5616     """
-> 5617     merge_result = dataset_update_method(self, other)
   5618     return self._replace(inplace=True, **merge_result._asdict())

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xarray/core/merge.py:1075, in dataset_update_method(dataset, other)
   1072             if coord_names:
   1073                 other[key] = value.drop_vars(coord_names)
-> 1075 return merge_core(
   1076     [dataset, other],
   1077     priority_arg=1,
   1078     indexes=dataset.xindexes,
   1079     combine_attrs="override",
   1080 )

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xarray/core/merge.py:711, in merge_core(objects, compat, join, combine_attrs, priority_arg, explicit_coords, indexes, fill_value, skip_align_args)
    708 skip_align_objs = [(pos, objects.pop(pos)) for pos in skip_align_args]
    710 coerced = coerce_pandas_values(objects)
--> 711 aligned = deep_align(
    712     coerced, join=join, copy=False, indexes=indexes, fill_value=fill_value
    713 )
    715 for pos, obj in skip_align_objs:
    716     aligned.insert(pos, obj)

File /global/u2/v/vo13/mambaforge/envs/e3sm_diags_dev_nompi_659/lib/python3.10/site-packages/xarray/core/alignment.py:946, in deep_align(objects, join, copy, indexes, exclude, raise_on_invalid, fill_value)
    943     else:
    944         out.append(variables)
--> 946 aligned = align(
    947     *targets,
    948     join=join,
    949     copy=copy,
    950     indexes=indexes,
    951     exclude=exclude,
...
    477         f"cannot reindex or align along dimension {dim!r} "
    478         f"because of conflicting dimension sizes: {sizes!r}" + add_err_msg
    479     )

ValueError: cannot reindex or align along dimension 'lat' because of conflicting dimension sizes: {1, 2} (note: an index is found along that dimension with size=1)
```


### Anything else we need to know?

_No response_

### Environment

Latest `main` xCDAT

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/xCDAT/xcdat/pull/653?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 100.00%. Comparing base [(`7745757`)](https://app.codecov.io/gh/xCDAT/xcdat/commit/77457577ebd358a86123049f78693e36719c8792?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT) to head [(`b7f6fe9`)](https://app.codecov.io/gh/xCDAT/xcdat/pull/653?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT).


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff            @@
##              main      #653   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           15        15           
  Lines         1543      1544    +1     
=========================================
+ Hits          1543      1544    +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/xCDAT/xcdat/pull/653?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=xCDAT).


Patch:
diff --git a/xcdat/regridder/regrid2.py b/xcdat/regridder/regrid2.py
index 823215d..9b71979 100644
--- a/xcdat/regridder/regrid2.py
+++ b/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)
 


Test patch:
diff --git a/tests/test_regrid.py b/tests/test_regrid.py
index c7253f4..d413cb0 100644
--- a/tests/test_regrid.py
+++ b/tests/test_regrid.py
@@ -372,6 +372,30 @@ class TestRegrid2Regridder:
             }
         )
 
+    # 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