Optimized Geometry Functions using Numba (#1072)

* re-write bounds to use Numba

* more numba

* fix all tests but two

* remove unused helper

* fix all tests but one

* recursion limit

* recursion limit

* cache and parallel

* fix failing test

* fix all tests but one

* fix final test

* Clean up commented out code

* cleanup comments

* more comment cleanup

* update use of face bounds

* fix failing test

* update docstrings (1)

* update API

* update docstrings (2)

* only use bounding box and raise warning for first computation

* add exception for non-face-centered data variables

* update docstring

* pre-commit

* update docstring and parameter

* add warning when translating numba functions for the first time

* update numba check

* remove print statement

Author: philipc2

docs/api.rst

@@ -126,6 +126,8 @@ Descriptors
    Grid.descriptors
    Grid.face_areas
    Grid.bounds
+   Grid.face_bounds_lon
+   Grid.face_bounds_lat
    Grid.edge_node_distances
    Grid.edge_face_distances
    Grid.antimeridian_face_indices

test/test_arcs.py

@@ -10,7 +10,7 @@
 import uxarray as ux
 
 from uxarray.grid.coordinates import _lonlat_rad_to_xyz
-from uxarray.grid.arcs import point_within_gca
+from uxarray.grid.arcs import point_within_gca, _point_within_gca_cartesian
 
 try:
     import constants
@@ -38,27 +38,27 @@ def test_pt_within_gcr(self):
 
         pt_same_lon_in = _lonlat_rad_to_xyz(0.0, 0.0)
         with self.assertRaises(ValueError):
-            point_within_gca(pt_same_lon_in, gcr_180degree_cart)
+            _point_within_gca_cartesian(pt_same_lon_in, gcr_180degree_cart)
 
         # Test when the point and the GCA all have the same longitude
         gcr_same_lon_cart = [
             _lonlat_rad_to_xyz(0.0, 1.5),
             _lonlat_rad_to_xyz(0.0, -1.5)
         ]
         pt_same_lon_in = _lonlat_rad_to_xyz(0.0, 0.0)
-        self.assertTrue(point_within_gca(pt_same_lon_in, gcr_same_lon_cart))
+        self.assertTrue(_point_within_gca_cartesian(pt_same_lon_in, gcr_same_lon_cart))
 
         pt_same_lon_out = _lonlat_rad_to_xyz(0.0, 1.500000000000001)
-        res = point_within_gca(pt_same_lon_out, gcr_same_lon_cart)
+        res = _point_within_gca_cartesian(pt_same_lon_out, gcr_same_lon_cart)
         self.assertFalse(res)
 
         pt_same_lon_out_2 = _lonlat_rad_to_xyz(0.1, 1.0)
-        res = point_within_gca(pt_same_lon_out_2, gcr_same_lon_cart)
+        res = _point_within_gca_cartesian(pt_same_lon_out_2, gcr_same_lon_cart)
         self.assertFalse(res)
 
         # And if we increase the digital place by one, it should be true again
         pt_same_lon_out_add_one_place = _lonlat_rad_to_xyz(0.0, 1.5000000000000001)
-        res = point_within_gca(pt_same_lon_out_add_one_place, gcr_same_lon_cart)
+        res = _point_within_gca_cartesian(pt_same_lon_out_add_one_place, gcr_same_lon_cart)
         self.assertTrue(res)
 
         # Normal case
@@ -69,7 +69,7 @@ def test_pt_within_gcr(self):
                                                         -0.997]])
         pt_cart_within = np.array(
             [0.25616109352676675, 0.9246590335292105, -0.010021496695000144])
-        self.assertTrue(point_within_gca(pt_cart_within, gcr_cart_2, True))
+        self.assertTrue(_point_within_gca_cartesian(pt_cart_within, gcr_cart_2, True))
 
         # Test other more complicate cases : The anti-meridian case
 
@@ -80,16 +80,16 @@ def test_pt_within_gcr_antimeridian(self):
         gcr_cart = np.array([[0.351, -0.724, 0.593], [0.617, 0.672, 0.410]])
         pt_cart = np.array(
             [0.9438777657502077, 0.1193199333436068, 0.922714737029319])
-        self.assertTrue(point_within_gca(pt_cart, gcr_cart, is_directed=True))
+        self.assertTrue(_point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=True))
         # If we swap the gcr, it should throw a value error since it's larger than 180 degree
         gcr_cart_flip = np.array([[0.617, 0.672, 0.410], [0.351, -0.724,
                                                           0.593]])
         with self.assertRaises(ValueError):
-            point_within_gca(pt_cart, gcr_cart_flip, is_directed=True)
+            _point_within_gca_cartesian(pt_cart, gcr_cart_flip, is_directed=True)
 
         # If we flip the gcr in the undirected mode, it should still work
         self.assertTrue(
-            point_within_gca(pt_cart, gcr_cart_flip, is_directed=False))
+            _point_within_gca_cartesian(pt_cart, gcr_cart_flip, is_directed=False))
 
         # 2nd anti-meridian case
         # GCR vertex0 in radian : [4.104711496596806, 0.5352983676533828],
@@ -100,9 +100,9 @@ def test_pt_within_gcr_antimeridian(self):
         pt_cart_within = np.array(
             [0.6136726305712109, 0.28442243941920053, -0.365605190899831])
         self.assertFalse(
-            point_within_gca(pt_cart_within, gcr_cart_1, is_directed=True))
+            _point_within_gca_cartesian(pt_cart_within, gcr_cart_1, is_directed=True))
         self.assertFalse(
-            point_within_gca(pt_cart_within, gcr_cart_1, is_directed=False))
+            _point_within_gca_cartesian(pt_cart_within, gcr_cart_1, is_directed=False))
 
         # The first case should not work and the second should work
         v1_rad = [0.1, 0.0]
@@ -112,10 +112,10 @@ def test_pt_within_gcr_antimeridian(self):
         gcr_cart = np.array([v1_cart, v2_cart])
         pt_cart = _lonlat_rad_to_xyz(0.01, 0.0)
         with self.assertRaises(ValueError):
-            point_within_gca(pt_cart, gcr_cart, is_directed=True)
+            _point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=True)
         gcr_car_flipped = np.array([v2_cart, v1_cart])
         self.assertTrue(
-            point_within_gca(pt_cart, gcr_car_flipped, is_directed=True))
+            _point_within_gca_cartesian(pt_cart, gcr_car_flipped, is_directed=True))
 
     def test_pt_within_gcr_cross_pole(self):
         gcr_cart = np.array([[0.351, 0.0, 0.3], [-0.351, 0.0, 0.3]])
@@ -125,7 +125,7 @@ def test_pt_within_gcr_cross_pole(self):
         # Normalize the point abd the GCA
         pt_cart = pt_cart / np.linalg.norm(pt_cart)
         gcr_cart = np.array([x / np.linalg.norm(x) for x in gcr_cart])
-        self.assertTrue(point_within_gca(pt_cart, gcr_cart, is_directed=False))
+        self.assertTrue(_point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=False))
 
         gcr_cart = np.array([[0.351, 0.0, 0.3], [-0.351, 0.0, -0.6]])
         pt_cart = np.array(
@@ -134,6 +134,6 @@ def test_pt_within_gcr_cross_pole(self):
         # When the point is not within the GCA
         pt_cart = pt_cart / np.linalg.norm(pt_cart)
         gcr_cart = np.array([x / np.linalg.norm(x) for x in gcr_cart])
-        self.assertFalse(point_within_gca(pt_cart, gcr_cart, is_directed=False))
+        self.assertFalse(_point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=False))
         with self.assertRaises(ValueError):
-            point_within_gca(pt_cart, gcr_cart, is_directed=True)
+            _point_within_gca_cartesian(pt_cart, gcr_cart, is_directed=True)

test/test_cross_sections.py

@@ -35,103 +35,98 @@ class TestQuadHex:
     All four faces intersect a constant latitude of 0.0
     """
 
-    @pytest.mark.parametrize("use_spherical_bounding_box", [True, False])
-    def test_constant_lat_cross_section_grid(self, use_spherical_bounding_box):
-
-
+    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, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+            uxgrid.cross_section.constant_latitude(lat=10.0, )
 
-    @pytest.mark.parametrize("use_spherical_bounding_box", [False])
-    def test_constant_lon_cross_section_grid(self, use_spherical_bounding_box):
+    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, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+        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)
 
-    @pytest.mark.parametrize("use_spherical_bounding_box", [False])
-    def test_constant_lat_cross_section_uxds(self, use_spherical_bounding_box):
+    def test_constant_lat_cross_section_uxds(self):
         uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path)
         uxds.uxgrid.normalize_cartesian_coordinates()
 
-        da_top_two = uxds['t2m'].cross_section.constant_latitude(lat=0.1, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+            uxds['t2m'].cross_section.constant_latitude(lat=10.0, )
+
 
-    @pytest.mark.parametrize("use_spherical_bounding_box", [False])
-    def test_constant_lon_cross_section_uxds(self, use_spherical_bounding_box):
+    def test_constant_lon_cross_section_uxds(self):
         uxds = ux.open_dataset(quad_hex_grid_path, quad_hex_data_path)
         uxds.uxgrid.normalize_cartesian_coordinates()
 
-        da_left_two = uxds['t2m'].cross_section.constant_longitude(lon=-0.1, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+        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, use_spherical_bounding_box=use_spherical_bounding_box)
+            uxds['t2m'].cross_section.constant_longitude(lon=10.0, )
 
 
 class TestCubeSphere:
-    @pytest.mark.parametrize("use_spherical_bounding_box", [True, False])
-    def test_north_pole(self, use_spherical_bounding_box):
+
+    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, use_spherical_bounding_box=use_spherical_bounding_box)
+            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
-    @pytest.mark.parametrize("use_spherical_bounding_box", [True, False])
-    def test_south_pole(self, use_spherical_bounding_box):
+
+    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, use_spherical_bounding_box=use_spherical_bounding_box)
+            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
 
@@ -152,8 +147,7 @@ def test_constant_lat(self):
 
         candidate_faces = constant_lat_intersections_face_bounds(
             lat=const_lat,
-            face_min_lat_rad=bounds_rad[:, 0, 0],
-            face_max_lat_rad=bounds_rad[:, 0, 1],
+            face_bounds_lat=bounds_rad[:, 0],
         )
 
         # Expected output
@@ -176,8 +170,7 @@ def test_constant_lat_out_of_bounds(self):
 
         candidate_faces = constant_lat_intersections_face_bounds(
             lat=const_lat,
-            face_min_lat_rad=bounds_rad[:, 0, 0],
-            face_max_lat_rad=bounds_rad[:, 0, 1],
+            face_bounds_lat=bounds_rad[:, 0],
         )
 
         assert len(candidate_faces) == 0