File: test_integration_zarr_conventions.py

package info (click to toggle)
python-rioxarray 0.22.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 7,572 kB
  • sloc: python: 8,937; makefile: 85
file content (211 lines) | stat: -rw-r--r-- 6,437 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
"""Integration tests for reading Zarr conventions."""
import numpy as np
import pyproj
import xarray as xr
from affine import Affine
from rasterio.crs import CRS

import rioxarray  # noqa: F401
from rioxarray import set_options
from rioxarray._convention import zarr
from rioxarray.enum import Convention


def _create_zarr_array_with_proj():
    """Create a DataArray with Zarr proj: convention attributes."""
    data = xr.DataArray(
        np.random.rand(10, 20),
        dims=["y", "x"],
        coords={
            "y": np.arange(10),
            "x": np.arange(20),
        },
    )
    data.attrs["zarr_conventions"] = [zarr.PROJ_CONVENTION]
    data.attrs["proj:wkt2"] = CRS.from_epsg(4326).to_wkt()
    return data


def _create_zarr_array_with_spatial():
    """Create a DataArray with Zarr spatial: convention attributes."""
    data = xr.DataArray(
        np.random.rand(10, 20),
        dims=["lat", "lon"],
        coords={
            "lat": np.arange(10),
            "lon": np.arange(20),
        },
    )
    data.attrs["zarr_conventions"] = [zarr.SPATIAL_CONVENTION]
    data.attrs["spatial:transform"] = [1.0, 0.0, 100.0, 0.0, -1.0, 200.0]
    data.attrs["spatial:dimensions"] = ["lat", "lon"]
    return data


def _create_zarr_array_with_both():
    """Create a DataArray with both Zarr conventions."""
    data = xr.DataArray(
        np.random.rand(10, 20),
        dims=["lat", "lon"],
        coords={
            "lat": np.arange(10),
            "lon": np.arange(20),
        },
    )
    data.attrs["zarr_conventions"] = [zarr.PROJ_CONVENTION, zarr.SPATIAL_CONVENTION]
    data.attrs["proj:wkt2"] = CRS.from_epsg(32618).to_wkt()
    data.attrs["spatial:transform"] = [10.0, 0.0, 500000.0, 0.0, -10.0, 4500000.0]
    data.attrs["spatial:dimensions"] = ["lat", "lon"]
    return data


def test_read_crs_from_zarr_convention():
    """Test reading CRS from DataArray with Zarr proj: convention."""
    data = _create_zarr_array_with_proj()

    crs = data.rio.crs
    assert crs is not None
    assert crs == CRS.from_epsg(4326)


def test_read_crs_from_zarr_convention__with_setting():
    """Test reading CRS with Convention.ZARR setting."""
    data = _create_zarr_array_with_proj()

    with set_options(convention=Convention.ZARR):
        crs = data.rio.crs
        assert crs is not None
        assert crs == CRS.from_epsg(4326)


def test_read_transform_from_zarr_convention():
    """Test reading transform from DataArray with Zarr spatial: convention."""
    data = _create_zarr_array_with_spatial()

    # Access transform via rio accessor
    # Check the cached version reads from Zarr spatial:transform
    cached = data.rio._cached_transform()
    assert cached is not None
    assert cached == Affine(1.0, 0.0, 100.0, 0.0, -1.0, 200.0)


def test_read_spatial_dimensions_from_zarr_convention():
    """Test reading spatial dimensions from Zarr spatial: convention."""
    data = _create_zarr_array_with_spatial()

    assert data.rio.x_dim == "lon"
    assert data.rio.y_dim == "lat"


def test_read_both_conventions():
    """Test reading from DataArray with both Zarr conventions."""
    data = _create_zarr_array_with_both()

    # CRS from proj:
    crs = data.rio.crs
    assert crs is not None
    assert crs == CRS.from_epsg(32618)

    # Transform from spatial:
    cached = data.rio._cached_transform()
    assert cached is not None
    assert cached == Affine(10.0, 0.0, 500000.0, 0.0, -10.0, 4500000.0)

    # Dimensions from spatial:
    assert data.rio.x_dim == "lon"
    assert data.rio.y_dim == "lat"


def test_fallback_zarr_to_cf():
    """Test that CF convention is tried as fallback when Zarr not found."""
    # Create data with CF convention
    data = xr.DataArray(
        np.random.rand(10, 20),
        dims=["y", "x"],
        coords={
            "y": np.arange(10),
            "x": np.arange(20),
        },
    )
    data.coords["spatial_ref"] = xr.Variable((), 0)
    data.coords["spatial_ref"].attrs["spatial_ref"] = "EPSG:4326"

    # Even with Zarr preference, should fall back to CF
    with set_options(convention=Convention.ZARR):
        crs = data.rio.crs
        assert crs is not None
        assert crs == CRS.from_epsg(4326)


def test_fallback_cf_to_zarr():
    """Test that Zarr convention is tried as fallback when CF not found."""
    # Create data with Zarr convention only
    data = _create_zarr_array_with_proj()

    # With CF preference (default), should fall back to Zarr
    crs = data.rio.crs
    assert crs is not None
    assert crs == CRS.from_epsg(4326)


def test_priority_zarr_over_cf():
    """Test that Zarr convention takes priority when setting is Zarr."""
    # Create data with both conventions (different CRS values)
    data = xr.DataArray(
        np.random.rand(10, 20),
        dims=["y", "x"],
        coords={
            "y": np.arange(10),
            "x": np.arange(20),
        },
    )
    # CF convention
    data.coords["spatial_ref"] = xr.Variable((), 0)
    data.coords["spatial_ref"].attrs["spatial_ref"] = "EPSG:4326"

    # Zarr convention (different CRS)
    data.attrs["zarr_conventions"] = [zarr.PROJ_CONVENTION]
    data.attrs["proj:wkt2"] = CRS.from_epsg(32618).to_wkt()

    # With Zarr setting, should prefer Zarr CRS
    with set_options(convention=Convention.ZARR):
        crs = data.rio.crs
        assert crs is not None
        assert crs == CRS.from_epsg(32618)

    # Reset to check default
    data2 = data.copy(deep=True)
    data2.rio._crs = None  # Reset cached CRS

    # With default setting (CF priority), should prefer CF CRS
    crs = data2.rio.crs
    assert crs is not None
    assert crs == CRS.from_epsg(4326)


def test_read_proj_code():
    """Test reading CRS from proj:code attribute."""
    data = xr.DataArray(
        np.random.rand(10, 20),
        dims=["y", "x"],
    )
    data.attrs["zarr_conventions"] = [zarr.PROJ_CONVENTION]
    data.attrs["proj:code"] = "EPSG:32618"

    crs = data.rio.crs
    assert crs is not None
    assert crs == CRS.from_epsg(32618)


def test_read_proj_projjson():
    """Test reading CRS from proj:projjson attribute."""
    data = xr.DataArray(
        np.random.rand(10, 20),
        dims=["y", "x"],
    )
    data.attrs["zarr_conventions"] = [zarr.PROJ_CONVENTION]
    data.attrs["proj:projjson"] = pyproj.CRS.from_epsg(4326).to_json_dict()

    crs = data.rio.crs
    assert crs is not None
    assert crs == CRS.from_epsg(4326)