Repository: zarr-developers/VirtualiZarr

Issue:
open_virtual_dataset returns some coordinates as data variables
The `xr.Dataset` constructed by `open_virtual_dataset` doesn't seem to correctly identify coordinates when the coordinate has more than one dimension. The bug seems to be in `separate_coords` [on this line](https://github.com/zarr-developers/VirtualiZarr/blob/f60707f86f6fa7415463c1a3e886eb134b84bdb0/virtualizarr/xarray.py#L363). The correct functionality could be to use the `coordinates` attribute within each variables `.zattrs` and maintain a set of all coordinate names

Here is a reproducible example:

```python
>>> import xarray as xr
>>> xr.tutorial.open_dataset("ROMS_example.nc") 
<xarray.Dataset> Size: 19MB
Dimensions:     (ocean_time: 2, s_rho: 30, eta_rho: 191, xi_rho: 371)
Coordinates:
    Cs_r        (s_rho) float64 240B ...
    lon_rho     (eta_rho, xi_rho) float64 567kB ...
    hc          float64 8B ...
    h           (eta_rho, xi_rho) float64 567kB ...
    lat_rho     (eta_rho, xi_rho) float64 567kB ...
    Vtransform  int32 4B ...
  * ocean_time  (ocean_time) datetime64[ns] 16B 2001-08-01 2001-08-08
  * s_rho       (s_rho) float64 240B -0.9833 -0.95 -0.9167 ... -0.05 -0.01667
Dimensions without coordinates: eta_rho, xi_rho
Data variables:
    salt        (ocean_time, s_rho, eta_rho, xi_rho) float32 17MB ...
    zeta        (ocean_time, eta_rho, xi_rho) float32 567kB ...
Attributes: (12/34)
    file:              ../output_20yr_obc/2001/ocean_his_0015.nc
    format:            netCDF-4/HDF5 file
    Conventions:       CF-1.4
    type:              ROMS/TOMS history file
    title:             TXLA ROMS hindcast run with dyes and oxygen
    rst_file:          ../output_20yr_obc/2001/ocean_rst.nc
    ...                ...
    compiler_flags:    -heap-arrays -fp-model fast -mt_mpi -ip -O3 -msse2 -free
    tiling:            010x012
    history:           Tue Jul 24 11:04:43 2018: /opt/nco/ncks -D 4 -t 8 /cop...
    ana_file:          /home/d.kobashi/TXLA_ROMS_reana/Functionals/ana_btflux...
    CPP_options:       TXLA2, ANA_BPFLUX, ANA_BSFLUX, ANA_BTFLUX, ANA_NUDGCOE...
    NCO:               netCDF Operators version 4.7.6-alpha04 (Homepage = htt...
```

```
% wget https://github.com/pydata/xarray-data/raw/master/ROMS_example.nc
```

```python
>>> from virtualizarr import open_virtual_dataset
>>> vds = open_virtual_dataset('ROMS_example.nc', indexes={})
>>> vds
<xarray.Dataset> Size: 19MB
Dimensions:     (ocean_time: 2, eta_rho: 191, xi_rho: 371, s_rho: 30)
Coordinates:
    s_rho       (s_rho) float64 240B ManifestArray<shape=(30,), dtype=float64...
    ocean_time  (ocean_time) float64 16B ManifestArray<shape=(2,), dtype=floa...
Dimensions without coordinates: eta_rho, xi_rho
Data variables:
    zeta        (ocean_time, eta_rho, xi_rho) float32 567kB ManifestArray<sha...
    lon_rho     (eta_rho, xi_rho) float64 567kB ManifestArray<shape=(191, 371...
    Vtransform  int32 4B ManifestArray<shape=(), dtype=int32, chunks=()>
    Cs_r        (s_rho) float64 240B ManifestArray<shape=(30,), dtype=float64...
    hc          float64 8B ManifestArray<shape=(), dtype=float64, chunks=()>
    lat_rho     (eta_rho, xi_rho) float64 567kB ManifestArray<shape=(191, 371...
    h           (eta_rho, xi_rho) float64 567kB ManifestArray<shape=(191, 371...
    salt        (ocean_time, s_rho, eta_rho, xi_rho) float32 17MB ManifestArr...
Attributes: (12/34)
    CPP_options:       TXLA2, ANA_BPFLUX, ANA_BSFLUX, ANA_BTFLUX, ANA_NUDGCOE...
    Conventions:       CF-1.4
    NCO:               netCDF Operators version 4.7.6-alpha04 (Homepage = htt...
    NLM_LBC:           \nEDGE:    WEST   SOUTH  EAST   NORTH  \nzeta:    Che ...
    ana_file:          /home/d.kobashi/TXLA_ROMS_reana/Functionals/ana_btflux...
    avg_base:          ../output_20yr_obc/2001/ocean_avg
    ...                ...
    sta_file:          ocean_sta.nc
    svn_rev:            
    svn_url:           https:://myroms.org/svn/src
    tiling:            010x012
    title:             TXLA ROMS hindcast run with dyes and oxygen
    type:              ROMS/TOMS history file
```

Note that the underlying kerchunk json does have this coordinate information since when you `virtualize` the dataset and materialize data, the coordinates are correct:

```python
>>> refs = vds.virtualize.to_kerchunk(filepath=None, format="dict")
>>> xr.open_dataset("reference://", engine="zarr", chunks={}, backend_kwargs={"storage_options": {"fo": refs, "consolidated": False}})
<xarray.Dataset> Size: 19MB
Dimensions:     (s_rho: 30, eta_rho: 191, xi_rho: 371, ocean_time: 2)
Coordinates:
    Cs_r        (s_rho) float64 240B dask.array<chunksize=(30,), meta=np.ndarray>
    Vtransform  float64 8B ...
    h           (eta_rho, xi_rho) float64 567kB dask.array<chunksize=(191, 371), meta=np.ndarray>
    hc          float64 8B ...
    lat_rho     (eta_rho, xi_rho) float64 567kB dask.array<chunksize=(191, 371), meta=np.ndarray>
    lon_rho     (eta_rho, xi_rho) float64 567kB dask.array<chunksize=(191, 371), meta=np.ndarray>
  * ocean_time  (ocean_time) datetime64[ns] 16B 2001-08-01 2001-08-08
  * s_rho       (s_rho) float64 240B -0.9833 -0.95 -0.9167 ... -0.05 -0.01667
Dimensions without coordinates: eta_rho, xi_rho
Data variables:
    salt        (ocean_time, s_rho, eta_rho, xi_rho) float32 17MB dask.array<chunksize=(1, 15, 96, 186), meta=np.ndarray>
    zeta        (ocean_time, eta_rho, xi_rho) float32 567kB dask.array<chunksize=(1, 191, 371), meta=np.ndarray>
Attributes: (12/34)
    CPP_options:       TXLA2, ANA_BPFLUX, ANA_BSFLUX, ANA_BTFLUX, ANA_NUDGCOE...
    Conventions:       CF-1.4
    NCO:               netCDF Operators version 4.7.6-alpha04 (Homepage = htt...
    NLM_LBC:           \nEDGE:    WEST   SOUTH  EAST   NORTH  \nzeta:    Che ...
    ana_file:          /home/d.kobashi/TXLA_ROMS_reana/Functionals/ana_btflux...
    avg_base:          ../output_20yr_obc/2001/ocean_avg
    ...                ...
    sta_file:          ocean_sta.nc
    svn_rev:            
    svn_url:           https:://myroms.org/svn/src
    tiling:            010x012
    title:             TXLA ROMS hindcast run with dyes and oxygen
    type:              ROMS/TOMS history file
```

Discussion:
TomNicholas: > I can refactor the code and update this PR to remove the coord_names param from separate_coords and simply use the newly added search within separate_coords

That would be great @ayushnag. I _think_ the presence of the `coord_names` kwarg comes from a time when I didn't actually understand where I was supposed to get the information about which variable was a coordinate, so I had left it general. But it seems we don't need it, so let's get rid of it.

We should also do a few other things:
- Strip the `coordinates` entry from the `.zattrs` - this is redundant information once we have set the variable as an xarray coordinate.
- Make sure the reverse process also works - saving a multi-dimensional coordinate to kerchunk json records the fact it should be a coordinate by re-adding an entry to the `.zattrs`.
- We need to test all that, so we should add another roundtripping test, which can use the `xr.tutorial.open_dataset("ROMS_example.nc")` you showed in #189. 
TomNicholas: @ayushnag I would like to merge this important bugfix and issue a release of a new version of virtualizarr. Are you likely to have time to come back to this or should I merge this PR and open another?
ayushnag: @TomNicholas I was looking into this some more and it seems that the `coord_names` param is needed for non dimension coordinate cases (there is coordinate information in `.zattrs`) and we can't eliminate that function param. When I tried to remove it, the [kerchunk roundtrip integration test failed](https://github.com/zarr-developers/VirtualiZarr/blob/179bb2ab42664546dd243a2e04fab8737846229a/virtualizarr/tests/test_integration.py#L148)

 I unfortunately don't have time to implement your new suggestions but you can feel free to merge this or make an updated PR. 
ayushnag: @TomNicholas I do have some time to add tests and get this merged. The current fix is in `common.py` however there is no `test_common`. Do you want me to make this file or just add the test to `test_backend` or `test_utils`? The test I am planning on adding is `test_separate_coords()` which will check the functionality of (1) and (2) from [your comment here](https://github.com/zarr-developers/VirtualiZarr/issues/281#issuecomment-2445526098)
TomNicholas: That would be great! Test_backends would be fine I think.
ayushnag: Great! Do you have an example netcdf that has some coordinates where CF decoding is required? It would be ideal to have a test that checks all three cases where two should pass after this PR and I can mark the third CF case as `xfail` with a TODO so it's a known issue/task for later. The problem with `xr.tutorial.open_dataset("ROMS_example")` is that it just has (1) and (2) so this PR actually finds all of the coordinates. 

If there isn't one to easily download I can manually construct one, I just need to know the cases where CF decoding is needed
TomNicholas: I think another one of the Xarray tutorial datasets (the roms one) might have this property. Does my half-finished PR to make (3) work include one?
ayushnag: I tested this change with ROMS_example and it does pass with just this PR
TomNicholas: Ah damn. Well we can just leave this for another PR.

If you do want to go further down the rabbit hole I think that one example of CF metadata that would cause Xarray to set extra variables as coordinates is when one variable has a `bounds` attribute that contains the name of another variable.

@sharkinsspatial maybe you know this bit of CF lore?

Patch:
diff --git a/virtualizarr/readers/common.py b/virtualizarr/readers/common.py
index f6f5dff..646d26c 100644
--- a/virtualizarr/readers/common.py
+++ b/virtualizarr/readers/common.py
@@ -144,8 +144,13 @@ def separate_coords(
     coord_vars: dict[
         str, tuple[Hashable, Any, dict[Any, Any], dict[Any, Any]] | Variable
     ] = {}
+    found_coord_names: set[str] = set()
+    # Search through variable attributes for coordinate names
+    for var in vars.values():
+        if "coordinates" in var.attrs:
+            found_coord_names.update(var.attrs["coordinates"].split(" "))
     for name, var in vars.items():
-        if name in coord_names or var.dims == (name,):
+        if name in coord_names or var.dims == (name,) or name in found_coord_names:
             # use workaround to avoid creating IndexVariables described here https://github.com/pydata/xarray/pull/8107#discussion_r1311214263
             if len(var.dims) == 1:
                 dim1d, *_ = var.dims


Test patch:
diff --git a/conftest.py b/conftest.py
index 810fd83..55c0782 100644
--- a/conftest.py
+++ b/conftest.py
@@ -35,6 +35,15 @@ def netcdf4_file(tmpdir):
     return filepath
 
 
+@pytest.fixture
+def netcdf4_file_with_2d_coords(tmpdir):
+    ds = xr.tutorial.open_dataset("ROMS_example")
+    filepath = f"{tmpdir}/ROMS_example.nc"
+    ds.to_netcdf(filepath, format="NETCDF4")
+    ds.close()
+    return filepath
+
+
 @pytest.fixture
 def netcdf4_virtual_dataset(netcdf4_file):
     from virtualizarr import open_virtual_dataset
diff --git a/virtualizarr/tests/test_backend.py b/virtualizarr/tests/test_backend.py
index 43a6bbd..e9b6081 100644
--- a/virtualizarr/tests/test_backend.py
+++ b/virtualizarr/tests/test_backend.py
@@ -156,6 +156,28 @@ class TestOpenVirtualDatasetAttrs:
         }
 
 
+@requires_kerchunk
+class TestDetermineCoords:
+    def test_infer_one_dimensional_coords(self, netcdf4_file):
+        vds = open_virtual_dataset(netcdf4_file, indexes={})
+        assert set(vds.coords) == {"time", "lat", "lon"}
+
+    def test_var_attr_coords(self, netcdf4_file_with_2d_coords):
+        vds = open_virtual_dataset(netcdf4_file_with_2d_coords, indexes={})
+
+        expected_dimension_coords = ["ocean_time", "s_rho"]
+        expected_2d_coords = ["lon_rho", "lat_rho", "h"]
+        expected_1d_non_dimension_coords = ["Cs_r"]
+        expected_scalar_coords = ["hc", "Vtransform"]
+        expected_coords = (
+            expected_dimension_coords
+            + expected_2d_coords
+            + expected_1d_non_dimension_coords
+            + expected_scalar_coords
+        )
+        assert set(vds.coords) == set(expected_coords)
+
+
 @network
 @requires_s3fs
 class TestReadFromS3: