Optimize Face Bounds Computation (#877)

* add benchmark for face bounds helpers * add benchmark for face bounds helpers * optimize point_within_gca * disable fma for benchmark * numba implementations of isclose and allclose * dot and cross numba * dot and cross numba * update dot * update dot * try to fix dot (1) * gca_gca_extreme * numba xyz to latlon rad with no norm * disable FMA * enable FMA * add more dot calls * remove leftover code * fix tests * fix failing tests * update isclose call * disable FMA for benchmark * fix broken isclose call * add docstrings, option to enable and disable fma * add some vectorization * numba all * point_within_gca numba * fix error in tests * better function name * revert asv.conf.json
UXARRAY · Aug 9, 2024 · 1a68daa · 1a68daa
1 parent 1790e0e
commit 1a68daa
Show file tree

Hide file tree

Showing 11 changed files with 308 additions and 144 deletions.
diff --git a/benchmarks/face_bounds.py b/benchmarks/face_bounds.py
@@ -10,8 +10,6 @@
 grid_scrip = current_path / "test" / "meshfiles" / "scrip" / "outCSne8" / "outCSne8.nc"
 grid_mpas= current_path / "test" / "meshfiles" / "mpas" / "QU" / "oQU480.231010.nc"
 
-
-
 class FaceBounds:
 
     params = [grid_quad_hex, grid_geoflow, grid_scrip, grid_mpas]

diff --git a/test/test_arcs.py b/test/test_arcs.py
@@ -36,7 +36,7 @@ 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(np.asarray(pt_same_lon_in), np.asarray(gcr_180degree_cart))
 
         gcr_180degree_cart = [
             _lonlat_rad_to_xyz(0.0, np.pi / 2.0),
@@ -45,27 +45,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(np.asarray(pt_same_lon_in), np.asarray(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(np.asarray(pt_same_lon_in), np.asarray(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(np.asarray(pt_same_lon_out), np.asarray(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(np.asarray(pt_same_lon_out_2), np.asarray(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(np.asarray(pt_same_lon_out_add_one_place), np.asarray(gcr_same_lon_cart))
         self.assertTrue(res)
 
         # Normal case
@@ -76,7 +76,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(np.asarray(pt_cart_within), np.asarray(gcr_cart_2), True))
 
         # Test other more complicate cases : The anti-meridian case
 
@@ -87,16 +87,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(np.asarray(pt_cart), np.asarray(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(np.asarray(pt_cart), np.asarray(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(np.asarray(pt_cart), np.asarray(gcr_cart_flip), is_directed=False))
 
         # 2nd anti-meridian case
         # GCR vertex0 in radian : [4.104711496596806, 0.5352983676533828],
@@ -107,9 +107,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(np.asarray(pt_cart_within), np.asarray(gcr_cart_1), is_directed=True))
         self.assertFalse(
-            point_within_gca(pt_cart_within, gcr_cart_1, is_directed=False))
+            point_within_gca(np.asarray(pt_cart_within), np.asarray(gcr_cart_1), is_directed=False))
 
         # The first case should not work and the second should work
         v1_rad = [0.1, 0.0]
@@ -119,10 +119,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(np.asarray(pt_cart), np.asarray(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(np.asarray(pt_cart), np.asarray(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]])

diff --git a/test/test_grid.py b/test/test_grid.py
@@ -952,8 +952,7 @@ def test_populate_bounds_GCA_mix(self):
         face_bounds = bounds_xarray.values
         nt.assert_allclose(grid.bounds.values, expected_bounds, atol=ERROR_TOLERANCE)
 
-    def test_populate_bounds_MPAS(self):
-        xrds = xr.open_dataset(self.gridfile_mpas)
-        uxgrid = ux.Grid.from_dataset(xrds, use_dual=True)
-        bounds_xarray = uxgrid.bounds
-        pass
+    def test_opti_bounds(self):
+        import uxarray
+        uxgrid = ux.open_grid(gridfile_CSne8)
+        bounds = uxgrid.bounds
diff --git a/uxarray/__init__.py b/uxarray/__init__.py
@@ -1,6 +1,8 @@
 # Sets the version of uxarray currently installeds
 # Attempt to import the needed modules
 
+import uxarray.constants
+
 from .core.api import open_grid, open_dataset, open_mfdataset
 
 from .core.dataset import UxDataset
@@ -21,6 +23,21 @@
     # Placeholder version incase an error occurs, such as the library isn't installed
     __version__ = "999"
 
+# Flag for enabling FMA instructions across the package
+
+
+def enable_fma():
+    """Enables Fused-Multiply-Add (FMA) instructions using the ``pyfma``
+    package."""
+    uxarray.constants.ENABLE_FMA = True
+
+
+def disable_fma():
+    """Disable Fused-Multiply-Add (FMA) instructions using the ``pyfma``
+    package."""
+    uxarray.constants.ENABLE_FMA = False
+
+
 __all__ = (
     "open_grid",
     "open_dataset",
@@ -30,4 +47,6 @@
     "INT_DTYPE",
     "INT_FILL_VALUE",
     "Grid",
+    "enable_fma",
+    "disable_fma",
 )
diff --git a/uxarray/constants.py b/uxarray/constants.py
@@ -15,4 +15,6 @@
 ENABLE_JIT_CACHE = True
 ENABLE_JIT = True
 
+ENABLE_FMA = False
+
 GRID_DIMS = ["n_node", "n_edge", "n_face"]