Fast Constant Latitude Cross Sections (#989)

* initial work on fast cross-sections

* add API

* add section for cross sections

* parallel numba

* update notebook

* add docstrings

* add benchmark

* update benchmark

* update quad hex grid

* add tests

* update tests

* update quad hex grid

* update name of intersection func

* update value error

* update user guide

* add docstring for uxda

* update docstring of intersections

* fix benchmark

* add tests for north and south pole

* Update intersections.py

* update user guide

* Delete docs/user-guide/grid.nc

* use accessor

* update notebook

* fix docs

Author: philipc2

benchmarks/mpas_ocean.py

@@ -4,6 +4,8 @@
 
 import uxarray as ux
 
+import numpy as np
+
 current_path = Path(os.path.dirname(os.path.realpath(__file__)))
 
 data_var = 'bottomDepth'
@@ -164,3 +166,11 @@ def teardown(self, resolution):
     def time_check_norm(self, resolution):
         from uxarray.grid.validation import _check_normalization
         _check_normalization(self.uxgrid)
+
+
+class CrossSections(DatasetBenchmark):
+    param_names = DatasetBenchmark.param_names + ['n_lat']
+    params = DatasetBenchmark.params + [[1, 2, 4, 8]]
+    def time_constant_lat_fast(self, resolution, n_lat):
+        for lat in np.linspace(-89, 89, n_lat):
+            self.uxds.uxgrid.constant_latitude_cross_section(lat, method='fast')

docs/api.rst

@@ -291,6 +291,30 @@ UxDataArray
    UxDataArray.subset.bounding_circle
 
 
+Cross Sections
+--------------
+
+
+Grid
+~~~~
+
+.. autosummary::
+   :toctree: generated/
+   :template: autosummary/accessor_method.rst
+
+   Grid.cross_section
+   Grid.cross_section.constant_latitude
+
+UxDataArray
+~~~~~~~~~~~
+
+.. autosummary::
+   :toctree: generated/
+   :template: autosummary/accessor_method.rst
+
+   UxDataArray.cross_section
+   UxDataArray.cross_section.constant_latitude
+
 Remapping
 ---------
 

docs/user-guide/cross-sections.ipynb

@@ -0,0 +1,208 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "4a432a8bf95d9cdb",
+   "metadata": {},
+   "source": [
+    "# Cross-Sections\n",
+    "\n",
+    "This section demonstrates how to extract cross-sections from an unstructured grid using UXarray, which allows the analysis and visualization across slices of grids.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b35ba4a2c30750e4",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-10-09T17:50:50.244285Z",
+     "start_time": "2024-10-09T17:50:50.239653Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import uxarray as ux\n",
+    "import geoviews.feature as gf\n",
+    "\n",
+    "import cartopy.crs as ccrs\n",
+    "import geoviews as gv\n",
+    "\n",
+    "projection = ccrs.Robinson()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "395a3db7-495c-4cff-b733-06bbe522a604",
+   "metadata": {},
+   "source": [
+    "## Data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b4160275c09fe6b0",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-10-09T17:50:51.217211Z",
+     "start_time": "2024-10-09T17:50:50.540946Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "base_path = \"../../test/meshfiles/ugrid/outCSne30/\"\n",
+    "grid_path = base_path + \"outCSne30.ug\"\n",
+    "data_path = base_path + \"outCSne30_vortex.nc\"\n",
+    "\n",
+    "uxds = ux.open_dataset(grid_path, data_path)\n",
+    "uxds[\"psi\"].plot(\n",
+    "    cmap=\"inferno\",\n",
+    "    periodic_elements=\"split\",\n",
+    "    projection=projection,\n",
+    "    title=\"Global Plot\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a7a40958-0a4d-47e4-9e38-31925261a892",
+   "metadata": {},
+   "source": [
+    "## Constant Latitude\n",
+    "\n",
+    "Cross-sections along constant latitude lines can be obtained using the ``.cross_section.constant_latitude`` method, available for both ``ux.Grid`` and ``ux.DataArray`` objects. This functionality allows users to extract and analyze slices of data at specified latitudes, providing insights into variations along horizontal sections of the grid.\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2fbe9f6e5bb59a17",
+   "metadata": {},
+   "source": [
+    "For example, we can obtain a cross-section at 30 degrees latitude by doing the following:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3775daa1-2f1d-4738-bab5-2b69ebd689d9",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-10-09T17:50:53.093314Z",
+     "start_time": "2024-10-09T17:50:53.077719Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "lat = 30\n",
+    "\n",
+    "uxda_constant_lat = uxds[\"psi\"].cross_section.constant_latitude(lat)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "dcec0b96b92e7f4",
+   "metadata": {},
+   "source": [
+    "Since the result is a new ``UxDataArray``, we can directly plot the result to see the cross-section."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "484b77a6-86da-4395-9e63-f5ac56e37deb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "(\n",
+    "    uxda_constant_lat.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 {lat} degrees latitude\",\n",
+    "    )\n",
+    "    * gf.grid(projection=projection)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7cca7de4722c121",
+   "metadata": {},
+   "source": [
+    "You can also perform operations on the cross-section, such as taking the mean."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1cbee722-34a4-4e67-8e22-f393d7d36c99",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(f\"Global Mean: {uxds['psi'].data.mean()}\")\n",
+    "print(f\"Mean at {lat} degrees lat: {uxda_constant_lat.data.mean()}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c4a7ee25-0b60-470f-bab7-92ff70563076",
+   "metadata": {},
+   "source": [
+    "## Constant Longitude"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9fcc8ec5-c6a8-4bde-a33d-7f37f9116ee2",
+   "metadata": {},
+   "source": [
+    "```{warning}\n",
+    "Constant longitude cross sections are not yet supported.\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "54d9eff1-67f1-4691-a3b0-1ee0c874c98f",
+   "metadata": {},
+   "source": [
+    "## Arbitrary Great Circle Arc (GCA)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ea94ff9f-fe86-470d-813b-45f32a633ffc",
+   "metadata": {},
+   "source": [
+    "```{warning}\n",
+    "Arbitrary great circle arc cross sections are not yet supported.\n",
+    "```"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/userguide.rst

@@ -43,6 +43,9 @@ These user guides provide detailed explanations of the core functionality in UXa
 `Subsetting <user-guide/subset.ipynb>`_
  Select specific regions of a grid
 
+`Cross-Sections <user-guide/cross-sections.ipynb>`_
+ Select cross-sections of a grid
+
 `Remapping <user-guide/remapping.ipynb>`_
  Remap (a.k.a Regrid) between unstructured grids
 
@@ -82,6 +85,7 @@ These user guides provide additional detail about specific features in UXarray.
    user-guide/mpl.ipynb
    user-guide/advanced-plotting.ipynb
    user-guide/subset.ipynb
+   user-guide/cross-sections.ipynb
    user-guide/remapping.ipynb
    user-guide/topological-aggregations.ipynb
    user-guide/calculus.ipynb

test/meshfiles/ugrid/quad-hexagon/grid.nc

@@ -0,0 +1,94 @@
+import uxarray as ux
+import pytest
+from pathlib import Path
+import os
+
+import numpy.testing as nt
+
+# Define the current path and file paths for grid and data
+current_path = Path(os.path.dirname(os.path.realpath(__file__)))
+quad_hex_grid_path = current_path / 'meshfiles' / "ugrid" / "quad-hexagon" / 'grid.nc'
+quad_hex_data_path = current_path / 'meshfiles' / "ugrid" / "quad-hexagon" / 'data.nc'
+
+cube_sphere_grid = current_path / "meshfiles" / "geos-cs" / "c12" / "test-c12.native.nc4"
+
+
+
+class TestQuadHex:
+    """The quad hexagon grid contains four faces.
+
+    Top Left Face: Index 1
+
+    Top Right Face: Index 2
+
+    Bottom Left Face: Index 0
+
+    Bottom Right Face: Index 3
+
+    The top two faces intersect a constant latitude of 0.1
+
+    The bottom two faces intersect a constant latitude of -0.1
+
+    All four faces intersect a constant latitude of 0.0
+    """
+
+    def test_constant_lat_cross_section_grid(self):
+        uxgrid = ux.open_grid(quad_hex_grid_path)
+
+        grid_top_two = uxgrid.cross_section.constant_latitude(lat=0.1)
+
+        assert grid_top_two.n_face == 2
+
+        grid_bottom_two = uxgrid.cross_section.constant_latitude(lat=-0.1)
+
+        assert grid_bottom_two.n_face == 2
+
+        grid_all_four = uxgrid.cross_section.constant_latitude(lat=0.0)
+
+        assert grid_all_four.n_face == 4
+
+        with pytest.raises(ValueError):
+            # no intersections found at this line
+            uxgrid.cross_section.constant_latitude(lat=10.0)
+
+
+    def test_constant_lat_cross_section_uxds(self):
+        uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path)
+
+        da_top_two = uxds['t2m'].cross_section.constant_latitude(lat=0.1)
+
+        nt.assert_array_equal(da_top_two.data, uxds['t2m'].isel(n_face=[1, 2]).data)
+
+        da_bottom_two = uxds['t2m'].cross_section.constant_latitude(lat=-0.1)
+
+        nt.assert_array_equal(da_bottom_two.data, uxds['t2m'].isel(n_face=[0, 3]).data)
+
+        da_all_four = uxds['t2m'].cross_section.constant_latitude(lat=0.0)
+
+        nt.assert_array_equal(da_all_four.data , uxds['t2m'].data)
+
+        with pytest.raises(ValueError):
+            # no intersections found at this line
+            uxds['t2m'].cross_section.constant_latitude(lat=10.0)
+
+
+class TestGeosCubeSphere:
+    def test_north_pole(self):
+        uxgrid = ux.open_grid(cube_sphere_grid)
+
+        lats = [89.85, 89.9, 89.95, 89.99]
+
+        for lat in lats:
+            cross_grid = uxgrid.cross_section.constant_latitude(lat=lat)
+            # Cube sphere grid should have 4 faces centered around the pole
+            assert cross_grid.n_face == 4
+
+    def test_south_pole(self):
+        uxgrid = ux.open_grid(cube_sphere_grid)
+
+        lats = [-89.85, -89.9, -89.95, -89.99]
+
+        for lat in lats:
+            cross_grid = uxgrid.cross_section.constant_latitude(lat=lat)
+            # Cube sphere grid should have 4 faces centered around the pole
+            assert cross_grid.n_face == 4

test/test_cross_sections.py

@@ -33,6 +33,7 @@
 from uxarray.plot.accessor import UxDataArrayPlotAccessor
 from uxarray.subset import DataArraySubsetAccessor
 from uxarray.remap import UxDataArrayRemapAccessor
+from uxarray.cross_sections import UxDataArrayCrossSectionAccessor
 from uxarray.core.aggregation import _uxda_grid_aggregate
 
 import warnings
@@ -85,6 +86,7 @@ def __init__(self, *args, uxgrid: Grid = None, **kwargs):
     plot = UncachedAccessor(UxDataArrayPlotAccessor)
     subset = UncachedAccessor(DataArraySubsetAccessor)
     remap = UncachedAccessor(UxDataArrayRemapAccessor)
+    cross_section = UncachedAccessor(UxDataArrayCrossSectionAccessor)
 
     def _repr_html_(self) -> str:
         if OPTIONS["display_style"] == "text":

uxarray/core/dataarray.py

@@ -0,0 +1,7 @@
+from .dataarray_accessor import UxDataArrayCrossSectionAccessor
+from .grid_accessor import GridCrossSectionAccessor
+
+__all__ = (
+    "GridCrossSectionAccessor",
+    "UxDataArrayCrossSectionAccessor",
+)