From 7ddd22f695a124a25e723cedcca34cd5bd6a2317 Mon Sep 17 00:00:00 2001 From: banesullivan Date: Sat, 2 Jul 2022 13:00:16 -0600 Subject: [PATCH 1/5] Add cf conventions helper --- pvxarray/accessor.py | 8 ++++++++ pvxarray/cf.py | 12 ++++++++++++ tests/test_cf.py | 17 +++++++++++++++++ 3 files changed, 37 insertions(+) create mode 100644 pvxarray/cf.py create mode 100644 tests/test_cf.py diff --git a/pvxarray/accessor.py b/pvxarray/accessor.py index 833afdc..188584e 100644 --- a/pvxarray/accessor.py +++ b/pvxarray/accessor.py @@ -4,6 +4,7 @@ import xarray as xr from pvxarray import rectilinear, structured +from pvxarray.cf import get_cf_names class _LocIndexer: @@ -51,6 +52,13 @@ def mesh( order: Optional[str] = None, component: Optional[str] = None, ) -> pv.DataSet: + if (3 - (x, y, z).count(None)) < 1: + try: + x, y, z = get_cf_names(self._obj) + except ImportError: + raise ValueError( + "You must specify at least one dimension as X, Y, or Z or install `cf_xarray`." + ) ndim = 0 if x is not None: _x = self._get_array(x) diff --git a/pvxarray/cf.py b/pvxarray/cf.py new file mode 100644 index 0000000..d6318ef --- /dev/null +++ b/pvxarray/cf.py @@ -0,0 +1,12 @@ +def get_cf_names(da): + """Use `cf_xarray` to get the names of the X, Y, and Z arrays.""" + try: + import cf_xarray # noqa + + axes = da.cf.axes + except (AttributeError, ImportError): + raise ImportError("Please import `cf_xarray` to use CF conventions.") + x = axes.get("X", [None])[0] + y = axes.get("Y", [None])[0] + z = axes.get("Z", [None])[0] + return x, y, z diff --git a/tests/test_cf.py b/tests/test_cf.py new file mode 100644 index 0000000..4b5be6d --- /dev/null +++ b/tests/test_cf.py @@ -0,0 +1,17 @@ +import numpy as np +import xarray as xr + + +def test_air_temperature(): + ds = xr.tutorial.load_dataset("air_temperature") + da = ds.air[dict(time=0)] + + mesh = da.pyvista.mesh() # No X,Y,Z specified, so should try cf_xarray + assert mesh + assert mesh.n_points == 1325 + assert "air" in mesh.point_data + + assert np.array_equal(mesh["air"], da.values.ravel()) + assert np.may_share_memory(mesh["air"], da.values.ravel()) + assert np.array_equal(mesh.x, da.lon) + assert np.array_equal(mesh.y, da.lat) From d0ff03abdf97e6f8517e6e637b24427d3dd41f86 Mon Sep 17 00:00:00 2001 From: banesullivan Date: Sat, 2 Jul 2022 13:04:20 -0600 Subject: [PATCH 2/5] add cf_xarray to depends --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index f07ff70..e56a8dd 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,3 +6,4 @@ netcdf4 pytest pytest-cov rioxarray +cf_xarray From c41080dc0b9e56f9fb6a0646b2b7fc7d73fa6323 Mon Sep 17 00:00:00 2001 From: banesullivan Date: Sat, 2 Jul 2022 13:06:51 -0600 Subject: [PATCH 3/5] Add time for future work --- pvxarray/accessor.py | 2 +- pvxarray/cf.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/pvxarray/accessor.py b/pvxarray/accessor.py index 188584e..9060e9f 100644 --- a/pvxarray/accessor.py +++ b/pvxarray/accessor.py @@ -54,7 +54,7 @@ def mesh( ) -> pv.DataSet: if (3 - (x, y, z).count(None)) < 1: try: - x, y, z = get_cf_names(self._obj) + x, y, z, _ = get_cf_names(self._obj) except ImportError: raise ValueError( "You must specify at least one dimension as X, Y, or Z or install `cf_xarray`." diff --git a/pvxarray/cf.py b/pvxarray/cf.py index d6318ef..9a25dcd 100644 --- a/pvxarray/cf.py +++ b/pvxarray/cf.py @@ -9,4 +9,5 @@ def get_cf_names(da): x = axes.get("X", [None])[0] y = axes.get("Y", [None])[0] z = axes.get("Z", [None])[0] - return x, y, z + t = axes.get("T", [None])[0] + return x, y, z, t From 1fb0a24fce486467d82cf8ee1d21951150694c23 Mon Sep 17 00:00:00 2001 From: banesullivan Date: Sat, 2 Jul 2022 13:08:50 -0600 Subject: [PATCH 4/5] Improve error message --- pvxarray/cf.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pvxarray/cf.py b/pvxarray/cf.py index 9a25dcd..6510c6d 100644 --- a/pvxarray/cf.py +++ b/pvxarray/cf.py @@ -5,7 +5,7 @@ def get_cf_names(da): axes = da.cf.axes except (AttributeError, ImportError): - raise ImportError("Please import `cf_xarray` to use CF conventions.") + raise ImportError("Please install `cf_xarray` to use CF conventions.") x = axes.get("X", [None])[0] y = axes.get("Y", [None])[0] z = axes.get("Z", [None])[0] From 949d7cc078b8501b58f82bd3341cf50f91e02d02 Mon Sep 17 00:00:00 2001 From: banesullivan Date: Sat, 2 Jul 2022 13:15:32 -0600 Subject: [PATCH 5/5] Clean up --- pvxarray/accessor.py | 10 ++++++---- pvxarray/cf.py | 2 +- tests/test_base.py | 5 +++++ 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/pvxarray/accessor.py b/pvxarray/accessor.py index 9060e9f..c836ed3 100644 --- a/pvxarray/accessor.py +++ b/pvxarray/accessor.py @@ -55,10 +55,12 @@ def mesh( if (3 - (x, y, z).count(None)) < 1: try: x, y, z, _ = get_cf_names(self._obj) - except ImportError: - raise ValueError( - "You must specify at least one dimension as X, Y, or Z or install `cf_xarray`." - ) + except ImportError: # pragma: no cover + pass + if (3 - (x, y, z).count(None)) < 1: + raise ValueError( + "You must specify at least one dimension as X, Y, or Z or install `cf_xarray`." + ) ndim = 0 if x is not None: _x = self._get_array(x) diff --git a/pvxarray/cf.py b/pvxarray/cf.py index 6510c6d..9f829cf 100644 --- a/pvxarray/cf.py +++ b/pvxarray/cf.py @@ -4,7 +4,7 @@ def get_cf_names(da): import cf_xarray # noqa axes = da.cf.axes - except (AttributeError, ImportError): + except (AttributeError, ImportError): # pragma: no cover raise ImportError("Please install `cf_xarray` to use CF conventions.") x = axes.get("X", [None])[0] y = axes.get("Y", [None])[0] diff --git a/tests/test_base.py b/tests/test_base.py index 13a1bdd..c5a0952 100644 --- a/tests/test_base.py +++ b/tests/test_base.py @@ -36,6 +36,11 @@ def test_report(): assert pvxarray.Report() +def test_no_cf_and_no_names(sample): + with pytest.raises(ValueError): + sample[dict(t=0)].pyvista.mesh() + + def test_bad_key(sample): with pytest.raises(KeyError): sample[dict(t=0)].pyvista.mesh(x="foo")