Skip to content

Commit 3fa5cff

Browse files
Search for coord_names in separate_coords (#191)
* find coord_names in vars * resolve merge conflict * add 2d coords test * add kerchunk dep and add 1d coord test --------- Co-authored-by: Tom Nicholas <tom@cworthy.org>
1 parent ab23caa commit 3fa5cff

File tree

3 files changed

+37
-1
lines changed

3 files changed

+37
-1
lines changed

conftest.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,15 @@ def netcdf4_file(tmpdir):
3535
return filepath
3636

3737

38+
@pytest.fixture
39+
def netcdf4_file_with_2d_coords(tmpdir):
40+
ds = xr.tutorial.open_dataset("ROMS_example")
41+
filepath = f"{tmpdir}/ROMS_example.nc"
42+
ds.to_netcdf(filepath, format="NETCDF4")
43+
ds.close()
44+
return filepath
45+
46+
3847
@pytest.fixture
3948
def netcdf4_virtual_dataset(netcdf4_file):
4049
from virtualizarr import open_virtual_dataset

virtualizarr/readers/common.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,8 +144,13 @@ def separate_coords(
144144
coord_vars: dict[
145145
str, tuple[Hashable, Any, dict[Any, Any], dict[Any, Any]] | Variable
146146
] = {}
147+
found_coord_names: set[str] = set()
148+
# Search through variable attributes for coordinate names
149+
for var in vars.values():
150+
if "coordinates" in var.attrs:
151+
found_coord_names.update(var.attrs["coordinates"].split(" "))
147152
for name, var in vars.items():
148-
if name in coord_names or var.dims == (name,):
153+
if name in coord_names or var.dims == (name,) or name in found_coord_names:
149154
# use workaround to avoid creating IndexVariables described here https://github.com/pydata/xarray/pull/8107#discussion_r1311214263
150155
if len(var.dims) == 1:
151156
dim1d, *_ = var.dims

virtualizarr/tests/test_backend.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,28 @@ def test_coordinate_variable_attrs_preserved(self, netcdf4_file):
156156
}
157157

158158

159+
@requires_kerchunk
160+
class TestDetermineCoords:
161+
def test_infer_one_dimensional_coords(self, netcdf4_file):
162+
vds = open_virtual_dataset(netcdf4_file, indexes={})
163+
assert set(vds.coords) == {"time", "lat", "lon"}
164+
165+
def test_var_attr_coords(self, netcdf4_file_with_2d_coords):
166+
vds = open_virtual_dataset(netcdf4_file_with_2d_coords, indexes={})
167+
168+
expected_dimension_coords = ["ocean_time", "s_rho"]
169+
expected_2d_coords = ["lon_rho", "lat_rho", "h"]
170+
expected_1d_non_dimension_coords = ["Cs_r"]
171+
expected_scalar_coords = ["hc", "Vtransform"]
172+
expected_coords = (
173+
expected_dimension_coords
174+
+ expected_2d_coords
175+
+ expected_1d_non_dimension_coords
176+
+ expected_scalar_coords
177+
)
178+
assert set(vds.coords) == set(expected_coords)
179+
180+
159181
@network
160182
@requires_s3fs
161183
class TestReadFromS3:

0 commit comments

Comments
 (0)