Skip to content

Commit

Permalink
Constant Longitude Cross Sections (#1011)
Browse files Browse the repository at this point in the history 
* add constant lon intersections

* add checks if edges are along a constant lat or lon

* add notebook

* add docstrings

* add docstring to numba func

* do not wrap around at poles for const lon

* Clean notebook

* add global error tollerance

---------

Co-authored-by: Rajeev Jain <[email protected]>
  • Loading branch information
@philipc2 @rajeeja
philipc2 and rajeeja authored Oct 31, 2024
1 parent 948b0f5 commit 191f4c1
Show file tree
Hide file tree
Showing 6 changed files with 330 additions and 15 deletions.
58 changes: 49 additions & 9 deletions docs/user-guide/cross-sections.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,10 @@
"id": "b35ba4a2c30750e4",
"metadata": {
"ExecuteTime": {
"end_time": "2024-10-09T17:50:50.244285Z",
"start_time": "2024-10-09T17:50:50.239653Z"
"start_time": "2024-10-14T16:39:35.957687Z"
},
"jupyter": {
"is_executing": true
}
},
"outputs": [],
Expand Down Expand Up @@ -64,6 +66,16 @@
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c844a3b5-ef71-41c8-a0ae-02ea937801d6",
"metadata": {},
"outputs": [],
"source": [
"uxds.uxgrid.normalize_cartesian_coordinates()"
]
},
{
"cell_type": "markdown",
"id": "a7a40958-0a4d-47e4-9e38-31925261a892",
Expand Down Expand Up @@ -94,7 +106,7 @@
},
"outputs": [],
"source": [
"lat = 30\n",
"lat = 0\n",
"\n",
"uxda_constant_lat = uxds[\"psi\"].cross_section.constant_latitude(lat)"
]
Expand Down Expand Up @@ -152,17 +164,45 @@
"id": "c4a7ee25-0b60-470f-bab7-92ff70563076",
"metadata": {},
"source": [
"## Constant Longitude"
"## Constant Longitude\n",
"\n",
"\n",
"\n",
"Cross-sections along constant longitude lines can be obtained using the ``.cross_section.constant_longitude`` method, available for both ``ux.Grid`` and ``ux.DataArray`` objects. \n"
]
},
{
"cell_type": "markdown",
"id": "9fcc8ec5-c6a8-4bde-a33d-7f37f9116ee2",
"cell_type": "code",
"execution_count": null,
"id": "f10917ce-568c-4e98-9b9c-d7c3c82e9ba3",
"metadata": {},
"outputs": [],
"source": [
"```{warning}\n",
"Constant longitude cross sections are not yet supported.\n",
"```"
"lon = 90\n",
"\n",
"uxda_constant_lon = uxds[\"psi\"].cross_section.constant_longitude(lon)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4add3a54-263e-41af-ac97-1e43c9141cb4",
"metadata": {},
"outputs": [],
"source": [
"(\n",
" uxda_constant_lon.plot(\n",
" rasterize=False,\n",
" backend=\"bokeh\",\n",
" cmap=\"inferno\",\n",
" projection=projection,\n",
" global_extent=True,\n",
" coastline=True,\n",
" title=f\"Cross Section at {lon} degrees longitude\",\n",
" periodic_elements=\"split\",\n",
" )\n",
" * gf.grid(projection=projection)\n",
")"
]
},
{
Expand Down
30 changes: 30 additions & 0 deletions test/test_cross_sections.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,21 @@ def test_constant_lat_cross_section_grid(self):
# no intersections found at this line
uxgrid.cross_section.constant_latitude(lat=10.0)

def test_constant_lon_cross_section_grid(self):
uxgrid = ux.open_grid(quad_hex_grid_path)

grid_left_two = uxgrid.cross_section.constant_longitude(lon=-0.1)

assert grid_left_two.n_face == 2

grid_right_two = uxgrid.cross_section.constant_longitude(lon=0.2)

assert grid_right_two.n_face == 2

with pytest.raises(ValueError):
# no intersections found at this line
uxgrid.cross_section.constant_longitude(lon=10.0)


def test_constant_lat_cross_section_uxds(self):
uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path)
Expand All @@ -71,6 +86,21 @@ def test_constant_lat_cross_section_uxds(self):
# no intersections found at this line
uxds['t2m'].cross_section.constant_latitude(lat=10.0)

def test_constant_lon_cross_section_uxds(self):
uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path)

da_left_two = uxds['t2m'].cross_section.constant_longitude(lon=-0.1)

nt.assert_array_equal(da_left_two.data, uxds['t2m'].isel(n_face=[0, 2]).data)

da_right_two = uxds['t2m'].cross_section.constant_longitude(lon=0.2)

nt.assert_array_equal(da_right_two.data, uxds['t2m'].isel(n_face=[1, 3]).data)

with pytest.raises(ValueError):
# no intersections found at this line
uxds['t2m'].cross_section.constant_longitude(lon=10.0)


class TestGeosCubeSphere:
def test_north_pole(self):
Expand Down
34 changes: 32 additions & 2 deletions uxarray/cross_sections/dataarray_accessor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,8 +54,38 @@ def constant_latitude(self, lat: float, method="fast"):

return self.uxda.isel(n_face=faces)

def constant_longitude(self, *args, **kwargs):
raise NotImplementedError
def constant_longitude(self, lon: float, method="fast"):
"""Extracts a cross-section of the data array at a specified constant
longitude.
Parameters
----------
lon : float
The longitude at which to extract the cross-section, in degrees.
method : str, optional
The internal method to use when identifying faces at the constant longitude.
Options are:
- 'fast': Uses a faster but potentially less accurate method for face identification.
- 'accurate': Uses a slower but more accurate method.
Default is 'fast'.
Raises
------
ValueError
If no intersections are found at the specified longitude, a ValueError is raised.
Examples
--------
>>> uxda.constant_longitude_cross_section(lon=-15.5)
Notes
-----
The accuracy and performance of the function can be controlled using the `method` parameter.
For higher precision requreiments, consider using method='acurate'.
"""
faces = self.uxda.uxgrid.get_faces_at_constant_longitude(lon, method)

return self.uxda.isel(n_face=faces)

def gca(self, *args, **kwargs):
raise NotImplementedError
Expand Down
60 changes: 58 additions & 2 deletions uxarray/cross_sections/grid_accessor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -81,8 +81,64 @@ def constant_latitude(self, lat: float, return_face_indices=False, method="fast"
else:
return grid_at_constant_lat

def constant_longitude(self, *args, **kwargs):
raise NotImplementedError
def constant_longitude(self, lon: float, return_face_indices=False, method="fast"):
"""Extracts a cross-section of the grid at a specified constant
longitude.
This method identifies and returns all faces (or grid elements) that intersect
with a given longitude. The returned cross-section can include either just the grid
or both the grid elements and the corresponding face indices, depending
on the `return_face_indices` parameter.
Parameters
----------
lon : float
The longitude at which to extract the cross-section, in degrees.
return_face_indices : bool, optional
If True, returns both the grid at the specified longitude and the indices
of the intersecting faces. If False, only the grid is returned.
Default is False.
method : str, optional
The internal method to use when identifying faces at the constant longitude.
Options are:
- 'fast': Uses a faster but potentially less accurate method for face identification.
- 'accurate': Uses a slower but more accurate method.
Default is 'fast'.
Returns
-------
grid_at_constant_lon : Grid
The grid with faces that interesected at a given longitudes
faces : array, optional
The indices of the faces that intersect with the specified longitude. This is only
returned if `return_face_indices` is set to True.
Raises
------
ValueError
If no intersections are found at the specified longitude, a ValueError is raised.
Examples
--------
>>> grid, indices = grid.cross_section.constant_longitude(lat=0.0, return_face_indices=True)
>>> grid = grid.cross_section.constant_longitude(lat=20.0)
Notes
-----
The accuracy and performance of the function can be controlled using the `method` parameter.
For higher precision requreiments, consider using method='acurate'.
"""
faces = self.uxgrid.get_faces_at_constant_longitude(lon, method)

if len(faces) == 0:
raise ValueError(f"No intersections found at lon={lon}.")

grid_at_constant_lon = self.uxgrid.isel(n_face=faces)

if return_face_indices:
return grid_at_constant_lon, faces
else:
return grid_at_constant_lon

def gca(self, *args, **kwargs):
raise NotImplementedError
Expand Down
100 changes: 100 additions & 0 deletions uxarray/grid/grid.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@

from uxarray.grid.intersections import (
fast_constant_lat_intersections,
fast_constant_lon_intersections,
)

from spatialpandas import GeoDataFrame
Expand Down Expand Up @@ -1059,6 +1060,40 @@ def edge_node_connectivity(self, value):
assert isinstance(value, xr.DataArray)
self._ds["edge_node_connectivity"] = value

@property
def edge_node_x(self) -> xr.DataArray:
"""Cartesian x location for the two nodes that make up every edge.
Dimensions: ``(n_edge, two)``
"""

if "edge_node_x" not in self._ds:
_edge_node_x = self.node_x.values[self.edge_node_connectivity.values]

self._ds["edge_node_x"] = xr.DataArray(
data=_edge_node_x,
dims=["n_edge", "two"],
)

return self._ds["edge_node_x"]

@property
def edge_node_y(self) -> xr.DataArray:
"""Cartesian y location for the two nodes that make up every edge.
Dimensions: ``(n_edge, two)``
"""

if "edge_node_y" not in self._ds:
_edge_node_y = self.node_y.values[self.edge_node_connectivity.values]

self._ds["edge_node_y"] = xr.DataArray(
data=_edge_node_y,
dims=["n_edge", "two"],
)

return self._ds["edge_node_y"]

@property
def edge_node_z(self) -> xr.DataArray:
"""Cartesian z location for the two nodes that make up every edge.
Expand Down Expand Up @@ -2084,3 +2119,68 @@ def get_faces_at_constant_latitude(self, lat, method="fast"):
faces = np.unique(self.edge_face_connectivity[edges].data.ravel())

return faces[faces != INT_FILL_VALUE]

def get_edges_at_constant_longitude(self, lon, method="fast"):
"""Identifies the edges of the grid that intersect with a specified
constant longitude.
This method computes the intersection of grid edges with a given longitude and
returns a collection of edges that cross or are aligned with that longitude.
The method used for identifying these edges can be controlled by the `method`
parameter.
Parameters
----------
lon : float
The longitude at which to identify intersecting edges, in degrees.
method : str, optional
The computational method used to determine edge intersections. Options are:
- 'fast': Uses a faster but potentially less accurate method for determining intersections.
- 'accurate': Uses a slower but more precise method.
Default is 'fast'.
Returns
-------
edges : array
A squeezed array of edges that intersect the specified constant longitude.
"""
if method == "fast":
edges = fast_constant_lon_intersections(
lon, self.edge_node_x.values, self.edge_node_y.values, self.n_edge
)
elif method == "accurate":
raise NotImplementedError("Accurate method not yet implemented.")
else:
raise ValueError(f"Invalid method: {method}.")
return edges.squeeze()

def get_faces_at_constant_longitude(self, lon, method="fast"):
"""Identifies the faces of the grid that intersect with a specified
constant longitude.
This method finds the faces (or cells) of the grid that intersect a given longitude
by first identifying the intersecting edges and then determining the faces connected
to these edges. The method used for identifying edges can be adjusted with the `method`
parameter.
Parameters
----------
lon : float
The longitude at which to identify intersecting faces, in degrees.
method : str, optional
The computational method used to determine intersecting edges. Options are:
- 'fast': Uses a faster but potentially less accurate method for determining intersections.
- 'accurate': Uses a slower but more precise method.
Default is 'fast'.
Returns
-------
faces : array
An array of unique face indices that intersect the specified longitude.
Faces that are invalid or missing (e.g., with a fill value) are excluded
from the result.
"""
edges = self.get_edges_at_constant_longitude(lon, method)
faces = np.unique(self.edge_face_connectivity[edges].data.ravel())

return faces[faces != INT_FILL_VALUE]
Loading

0 comments on commit 191f4c1

Please sign in to comment.