diff --git a/docs/api.rst b/docs/api.rst index e91591c2e..dd89363ba 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -250,6 +250,8 @@ Grid Accessor :toctree: generated/ UxDataArray.uxgrid + UxDataArray.data_mapping + UxDataArray.data_location UxDataset diff --git a/test/core/test_dataarray.py b/test/core/test_dataarray.py index 0e0d19aaa..b961a5a37 100644 --- a/test/core/test_dataarray.py +++ b/test/core/test_dataarray.py @@ -134,3 +134,32 @@ def test_isel_invalid_dim(gridpath, datasetpath): match=r"Dimensions \{'level'\} do not exist\..*Available dimensions: \('time', 'n_face'\)", ): uxda.isel(level=0) + + +def test_data_location(): + """Tests data_location for face/node/edge centered data and non-grid data.""" + uxgrid = ux.Grid.from_healpix(zoom=1) + + face_da = UxDataArray( + np.ones(uxgrid.n_face), dims=["n_face"], uxgrid=uxgrid + ) + node_da = UxDataArray( + np.ones(uxgrid.n_node), dims=["n_node"], uxgrid=uxgrid + ) + edge_da = UxDataArray( + np.ones(uxgrid.n_edge), dims=["n_edge"], uxgrid=uxgrid + ) + other_da = UxDataArray( + np.ones(5), dims=["other_dim"], uxgrid=uxgrid + ) + + assert face_da.data_location == "face_centered" + assert node_da.data_location == "node_centered" + assert edge_da.data_location == "edge_centered" + assert other_da.data_location is None + + # Works when an extra (non-grid) dimension is present + face_time = UxDataArray( + np.ones((3, uxgrid.n_face)), dims=["time", "n_face"], uxgrid=uxgrid + ) + assert face_time.data_location == "face_centered" diff --git a/uxarray/core/dataarray.py b/uxarray/core/dataarray.py index 478e3c2e3..b2516e4e4 100644 --- a/uxarray/core/dataarray.py +++ b/uxarray/core/dataarray.py @@ -164,6 +164,40 @@ def data_mapping(self): else: return None + @property + def data_location(self): + """Returns where on the grid the data variable is stored. + + The location is inferred from the grid dimension present in the data + variable, using UGRID-style names: + + - ``"face_centered"`` if the data contains the ``n_face`` dimension + - ``"node_centered"`` if the data contains the ``n_node`` dimension + - ``"edge_centered"`` if the data contains the ``n_edge`` dimension + - ``None`` if the data is not mapped to the grid + + Notes + ----- + Additional locations described in the UGRID/spectral-element ecosystem + (e.g. ``"face_average"``, ``"edge_orthogonal"``, ``"edge_parallel"``, + ``"cgll"``, ``"dgll"``) cannot be inferred from a data variable's + dimensions alone and are not currently distinguished here. + + Returns + ------- + str or None + One of ``"face_centered"``, ``"node_centered"``, + ``"edge_centered"``, or ``None``. + """ + if self._face_centered(): + return "face_centered" + elif self._node_centered(): + return "node_centered" + elif self._edge_centered(): + return "edge_centered" + else: + return None + def to_geodataframe( self, periodic_elements: str | None = "exclude",