Skip to content

Commit 630d40d

Browse files
aulemahaldcherian
aulemahal
and
dcherian
authored
Compute vertices from bounds (#108)
* Add cf.get_corners - tests - test dataset * Remove leftover code and update whats new * blacken * fix whatsnew * Move to new helpers file - add order option * run precommit * add new file! * Forgot new test file * Apply suggestions from code review * Minor updates. * Change corners to vertices * One forgotten word Co-authored-by: dcherian <[email protected]>
1 parent 5c8844c commit 630d40d

File tree

9 files changed

+408
-1
lines changed

9 files changed

+408
-1
lines changed

‎cf_xarray/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
from .accessor import CFAccessor # noqa
2+
from .helpers import bounds_to_vertices, vertices_to_bounds # noqa

‎cf_xarray/accessor.py

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
import xarray as xr
2121
from xarray import DataArray, Dataset
2222

23+
from .helpers import bounds_to_vertices
24+
2325
#: Classes wrapped by cf_xarray.
2426
_WRAPPED_CLASSES = (
2527
xr.core.resample.Resample,
@@ -1328,6 +1330,87 @@ def add_bounds(self, dims: Union[Hashable, Iterable[Hashable]]):
13281330

13291331
return self._maybe_to_dataarray(obj)
13301332

1333+
def bounds_to_vertices(
1334+
self,
1335+
keys: Union[str, Iterable[str]] = None,
1336+
order: str = "counterclockwise",
1337+
) -> Dataset:
1338+
"""
1339+
Convert bounds variable to vertices.
1340+
1341+
There 2 covered cases:
1342+
- 1D coordinates, with bounds of shape (N, 2),
1343+
converted to vertices of shape (N+1,)
1344+
- 2D coordinates, with bounds of shape (N, M, 4).
1345+
converted to vertices of shape (N+1, M+1).
1346+
1347+
Parameters
1348+
----------
1349+
keys : str or Iterable[str], optional
1350+
The names of the variables whose bounds are to be converted to vertices.
1351+
If not given, converts all available bounds within self.cf.keys().
1352+
order : {'counterclockwise', 'clockwise', None}
1353+
Valid for 2D coordinates only (bounds of shape (N, M, 4), ignored otherwise.
1354+
Order the bounds are given in, assuming that ax0-ax1-upward is a right
1355+
handed coordinate system, where ax0 and ax1 are the two first dimensions of
1356+
the variable. If None, the counterclockwise version is computed and then
1357+
verified. If the check fails the clockwise version is returned.
1358+
1359+
Returns
1360+
-------
1361+
Dataset
1362+
Copy of the dataset with added vertices variables.
1363+
Either of shape (N+1,) or (N+1, M+1). New vertex dimensions are named
1364+
from the intial dimension and suffix "_vertices". Variables with similar
1365+
names are overwritten.
1366+
1367+
Raises
1368+
------
1369+
ValueError
1370+
If any of the keys given doesn't corresponds to existing bounds.
1371+
1372+
Notes
1373+
-----
1374+
Getting the correct axes "order" is tricky. There are no real standards for
1375+
dimension names or even axes order, even though the CF conventions mentions the
1376+
ax0-ax1-upward (counterclockwise bounds) as being the default. Moreover, xarray can
1377+
tranpose data without raising any warning or error, which make attributes
1378+
unreliable.
1379+
1380+
Please refer to the CF conventions document : http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#cell-boundaries.
1381+
"""
1382+
if keys is None:
1383+
coords: Iterable[str] = self.keys()
1384+
elif isinstance(keys, str):
1385+
coords = (keys,)
1386+
else:
1387+
coords = keys
1388+
1389+
obj = self._maybe_to_dataset(self._obj.copy(deep=True))
1390+
1391+
for coord in coords:
1392+
try:
1393+
bounds = self.get_bounds(coord)
1394+
except KeyError as exc:
1395+
if keys is not None:
1396+
raise ValueError(
1397+
f"vertices are computed from bounds but given key {coord} did not correspond to existing bounds."
1398+
) from exc
1399+
else:
1400+
name = f"{self[coord].name}_vertices"
1401+
obj = obj.assign( # Overwrite any variable with the same name.
1402+
{
1403+
name: bounds_to_vertices(
1404+
bounds,
1405+
bounds_dim=list(set(bounds.dims) - set(self[coord].dims))[
1406+
0
1407+
],
1408+
order=order,
1409+
)
1410+
}
1411+
)
1412+
return obj
1413+
13311414
def decode_vertical_coords(self, prefix="z"):
13321415
"""
13331416
Decode parameterized vertical coordinates in place.

‎cf_xarray/helpers.py

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
from typing import Optional, Sequence
2+
3+
import numpy as np
4+
import xarray as xr
5+
from xarray import DataArray
6+
7+
8+
def bounds_to_vertices(
9+
bounds: DataArray, bounds_dim: str, order: Optional[str] = "counterclockwise"
10+
) -> DataArray:
11+
"""
12+
Convert bounds variable to vertices. There 2 covered cases:
13+
- 1D coordinates, with bounds of shape (N, 2),
14+
converted to vertices of shape (N+1,)
15+
- 2D coordinates, with bounds of shape (N, M, 4).
16+
converted to vertices of shape (N+1, M+1).
17+
18+
Parameters
19+
----------
20+
bounds: DataArray
21+
The bounds to convert. Must be of shape (N, 2) or (N, M, 4).
22+
bounds_dim : str
23+
The name of the bounds dimension of `bounds` (the one of length 2 or 4).
24+
order : {'counterclockwise', 'clockwise', None}
25+
Valid for 2D coordinates only (bounds of shape (N, M, 4), ignored otherwise.
26+
Order the bounds are given in, assuming that ax0-ax1-upward is a right handed
27+
coordinate system, where ax0 and ax1 are the two first dimensions of `bounds`.
28+
If None, the counterclockwise version is computed and then verified. If the
29+
check fails the clockwise version is returned. See Notes for more details.
30+
31+
Returns
32+
-------
33+
DataArray
34+
Either of shape (N+1,) or (N+1, M+1). New vertex dimensions are named
35+
from the intial dimension and suffix "_vertices".
36+
37+
Notes
38+
-----
39+
Getting the correct axes "order" is tricky. There are no real standards for
40+
dimension names or even axes order, even though the CF conventions mentions the
41+
ax0-ax1-upward (counterclockwise bounds) as being the default. Moreover, xarray can
42+
tranpose data without raising any warning or error, which make attributes
43+
unreliable.
44+
45+
Please refer to the CF conventions document : http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#cell-boundaries.
46+
"""
47+
# Get old and new dimension names and retranspose array to have bounds dim at axis 0.
48+
bnd_dim = (
49+
bounds_dim if isinstance(bounds_dim, str) else bounds.get_axis_num(bounds_dim)
50+
)
51+
old_dims = [dim for dim in bounds.dims if dim != bnd_dim]
52+
new_dims = [f"{dim}_vertices" for dim in old_dims]
53+
values = bounds.transpose(bnd_dim, *old_dims).data
54+
if len(old_dims) == 2 and bounds.ndim == 3 and bounds[bnd_dim].size == 4:
55+
# Vertices case (2D lat/lon)
56+
if order in ["counterclockwise", None]:
57+
# Names assume we are drawing axis 1 upward et axis 2 rightward.
58+
bot_left = values[0, :, :]
59+
bot_right = values[1, :, -1:]
60+
top_right = values[2, -1:, -1:]
61+
top_left = values[3, -1:, :]
62+
vertex_vals = np.block([[bot_left, bot_right], [top_left, top_right]])
63+
if order is None: # We verify if the ccw version works.
64+
calc_bnds = vertices_to_bounds(vertex_vals).values
65+
order = "counterclockwise" if np.all(calc_bnds == values) else "clockwise"
66+
if order == "clockwise":
67+
bot_left = values[0, :, :]
68+
top_left = values[1, -1:, :]
69+
top_right = values[2, -1:, -1:]
70+
bot_right = values[3, :, -1:]
71+
# Our asumption was wrong, axis 1 is rightward and axis 2 is upward
72+
vertex_vals = np.block([[bot_left, bot_right], [top_left, top_right]])
73+
elif len(old_dims) == 1 and bounds.ndim == 2 and bounds[bnd_dim].size == 2:
74+
# Middle points case (1D lat/lon)
75+
vertex_vals = np.concatenate((values[0, :], values[1, -1:]))
76+
else:
77+
raise ValueError(
78+
f"Bounds format not understood. Got {bounds.dims} with shape {bounds.shape}."
79+
)
80+
81+
return xr.DataArray(vertex_vals, dims=new_dims)
82+
83+
84+
def vertices_to_bounds(
85+
vertices: DataArray, out_dims: Sequence[str] = ("bounds", "x", "y")
86+
) -> DataArray:
87+
"""
88+
Convert vertices to CF-compliant bounds. There 2 covered cases:
89+
- 1D coordinates, with vertices of shape (N+1,),
90+
converted to bounds of shape (N, 2)
91+
- 2D coordinates, with vertices of shape (N+1, M+1).
92+
converted to bounds of shape (N, M, 4).
93+
94+
Parameters
95+
----------
96+
bounds: DataArray
97+
The bounds to convert. Must be of shape (N, 2) or (N, M, 4).
98+
out_dims : Sequence[str],
99+
The name of the dimension in the output. The first is the 'bounds'
100+
dimension and the following are the coordinate dimensions.
101+
Returns
102+
-------
103+
DataArray
104+
"""
105+
if vertices.ndim == 1:
106+
bnd_vals = np.stack((vertices[:-1], vertices[1:]), axis=0)
107+
elif vertices.ndim == 2:
108+
bnd_vals = np.stack(
109+
(
110+
vertices[:-1, :-1],
111+
vertices[:-1, 1:],
112+
vertices[1:, 1:],
113+
vertices[1:, :-1],
114+
),
115+
axis=0,
116+
)
117+
else:
118+
raise ValueError(
119+
f"vertices format not understood. Got {vertices.dims} with shape {vertices.shape}."
120+
)
121+
return xr.DataArray(bnd_vals, dims=out_dims[: vertices.ndim + 1])

‎cf_xarray/tests/datasets.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -130,3 +130,61 @@
130130
np.random.randn(2, 30),
131131
{"positive": "up"},
132132
)
133+
134+
135+
# Dataset with random data on a grid that is some sort of Mollweide projection
136+
XX, YY = np.mgrid[:11, :11] * 5 - 25
137+
XX_bnds, YY_bnds = np.mgrid[:12, :12] * 5 - 27.5
138+
139+
R = 50
140+
theta = np.arcsin(YY / (R * np.sqrt(2)))
141+
lat = np.rad2deg(np.arcsin((2 * theta + np.sin(2 * theta)) / np.pi))
142+
lon = np.rad2deg(XX * np.pi / (R * 2 * np.sqrt(2) * np.cos(theta)))
143+
144+
theta_bnds = np.arcsin(YY_bnds / (R * np.sqrt(2)))
145+
lat_vertices = np.rad2deg(np.arcsin((2 * theta_bnds + np.sin(2 * theta_bnds)) / np.pi))
146+
lon_vertices = np.rad2deg(XX_bnds * np.pi / (R * 2 * np.sqrt(2) * np.cos(theta_bnds)))
147+
148+
lon_bounds = np.stack(
149+
(
150+
lon_vertices[:-1, :-1],
151+
lon_vertices[:-1, 1:],
152+
lon_vertices[1:, 1:],
153+
lon_vertices[1:, :-1],
154+
),
155+
axis=0,
156+
)
157+
lat_bounds = np.stack(
158+
(
159+
lat_vertices[:-1, :-1],
160+
lat_vertices[:-1, 1:],
161+
lat_vertices[1:, 1:],
162+
lat_vertices[1:, :-1],
163+
),
164+
axis=0,
165+
)
166+
167+
mollwds = xr.Dataset(
168+
coords=dict(
169+
lon=xr.DataArray(
170+
lon,
171+
dims=("x", "y"),
172+
attrs={"units": "degrees_east", "bounds": "lon_bounds"},
173+
),
174+
lat=xr.DataArray(
175+
lat,
176+
dims=("x", "y"),
177+
attrs={"units": "degrees_north", "bounds": "lat_bounds"},
178+
),
179+
),
180+
data_vars=dict(
181+
lon_bounds=xr.DataArray(
182+
lon_bounds, dims=("bounds", "x", "y"), attrs={"units": "degrees_east"}
183+
),
184+
lat_bounds=xr.DataArray(
185+
lat_bounds, dims=("bounds", "x", "y"), attrs={"units": "degrees_north"}
186+
),
187+
lon_vertices=xr.DataArray(lon_vertices, dims=("x_vertices", "y_vertices")),
188+
lat_vertices=xr.DataArray(lat_vertices, dims=("x_vertices", "y_vertices")),
189+
),
190+
)

‎cf_xarray/tests/test_accessor.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -429,6 +429,32 @@ def test_bounds():
429429
assert_identical(actual, expected)
430430

431431

432+
def test_bounds_to_vertices():
433+
# All available
434+
ds = airds.cf.add_bounds(["lon", "lat"])
435+
dsc = ds.cf.bounds_to_vertices()
436+
assert "lon_vertices" in dsc
437+
assert "lat_vertices" in dsc
438+
439+
# Giving key
440+
dsc = ds.cf.bounds_to_vertices("longitude")
441+
assert "lon_vertices" in dsc
442+
assert "lat_vertices" not in dsc
443+
444+
dsc = ds.cf.bounds_to_vertices(["longitude", "latitude"])
445+
assert "lon_vertices" in dsc
446+
assert "lat_vertices" in dsc
447+
448+
# Error
449+
with pytest.raises(ValueError):
450+
dsc = ds.cf.bounds_to_vertices("T")
451+
452+
# Words on datetime arrays to
453+
ds = airds.cf.add_bounds("time")
454+
dsc = ds.cf.bounds_to_vertices()
455+
assert "time_bounds" in dsc
456+
457+
432458
def test_docstring():
433459
assert "One of ('X'" in airds.cf.groupby.__doc__
434460
assert "One or more of ('X'" in airds.cf.mean.__doc__

‎cf_xarray/tests/test_helpers.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
from numpy.testing import assert_array_equal
2+
from xarray.testing import assert_equal
3+
4+
import cf_xarray as cfxr # noqa
5+
6+
from .datasets import airds, mollwds
7+
8+
try:
9+
from dask.array import Array as DaskArray
10+
except ImportError:
11+
DaskArray = None
12+
13+
14+
def test_bounds_to_vertices():
15+
# 1D case
16+
ds = airds.cf.add_bounds(["lon", "lat", "time"])
17+
lat_c = cfxr.bounds_to_vertices(ds.lat_bounds, bounds_dim="bounds")
18+
assert_array_equal(ds.lat.values + 1.25, lat_c.values[:-1])
19+
20+
# 2D case, CF- order
21+
lat_c = cfxr.bounds_to_vertices(mollwds.lat_bounds, bounds_dim="bounds")
22+
assert_equal(mollwds.lat_vertices, lat_c)
23+
24+
# Transposing the array changes the bounds direction
25+
ds = mollwds.transpose("bounds", "y", "x", "y_vertices", "x_vertices")
26+
lon_c = cfxr.bounds_to_vertices(
27+
ds.lon_bounds, bounds_dim="bounds", order="clockwise"
28+
)
29+
lon_c2 = cfxr.bounds_to_vertices(ds.lon_bounds, bounds_dim="bounds", order=None)
30+
assert_equal(ds.lon_vertices, lon_c)
31+
assert_equal(ds.lon_vertices, lon_c2)
32+
33+
# Preserves dask-backed arrays
34+
if DaskArray is not None:
35+
lon_bounds = ds.lon_bounds.chunk()
36+
lon_c = cfxr.bounds_to_vertices(
37+
lon_bounds, bounds_dim="bounds", order="clockwise"
38+
)
39+
assert isinstance(lon_c.data, DaskArray)
40+
41+
42+
def test_vertices_to_bounds():
43+
# 1D case
44+
ds = airds.cf.add_bounds(["lon", "lat", "time"])
45+
lat_c = cfxr.bounds_to_vertices(ds.lat_bounds, bounds_dim="bounds")
46+
lat_b = cfxr.vertices_to_bounds(lat_c, out_dims=("bounds", "lat"))
47+
assert_array_equal(ds.lat_bounds, lat_b)
48+
49+
# Datetime
50+
time_c = cfxr.bounds_to_vertices(ds.time_bounds, bounds_dim="bounds")
51+
time_b = cfxr.vertices_to_bounds(time_c, out_dims=("bounds", "time"))
52+
assert_array_equal(ds.time_bounds, time_b)
53+
54+
# 2D case
55+
lon_b = cfxr.vertices_to_bounds(mollwds.lon_vertices, out_dims=("bounds", "x", "y"))
56+
assert_array_equal(mollwds.lon_bounds, lon_b)

‎doc/api.rst

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,9 +40,21 @@ Dataset
4040
Dataset.cf.decode_vertical_coords
4141
Dataset.cf.describe
4242
Dataset.cf.get_bounds
43+
Dataset.cf.bounds_to_vertices
4344
Dataset.cf.standard_names
4445
Dataset.cf.keys
4546
Dataset.cf.axes
4647
Dataset.cf.coordinates
4748
Dataset.cf.guess_coord_axis
4849
Dataset.cf.rename_like
50+
51+
.. currentmodule:: cf_xarray
52+
53+
Top-level API
54+
-------------
55+
56+
.. autosummary::
57+
:toctree: generated/
58+
59+
bounds_to_vertices
60+
vertices_to_bounds

0 commit comments

Comments
 (0)